#!/usr/bin/env python3
"""
Pilot Runner v4 - Critical Fixes for Baseline Comparison

Fixes from v3:
1. Chinese short reveal detection (简短/字以内/不超过 etc.)
2. Symmetric bullets constraint (has_bullets violation for require_bullets=False)
3. Better wrong_lang with CJK ratio + math exemption
4. Persona-conditional query templates (no self-contradiction)
5. Violation-triggered complaint mechanism (for online RL signal)

This version is ready for proper baseline comparison.
"""

import sys
import os
import re
import json
from datetime import datetime
from dataclasses import dataclass, asdict, field
from typing import List, Dict, Any, Optional, Tuple, Set, Literal

# Add src to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../src"))

from personalization.serving import PersonalizedLLM, Feedback, AssistantResponse


# =============================================================================
# Style Preferences (True Preferences)
# =============================================================================

@dataclass
class StylePrefs:
    """User's TRUE style preferences."""
    require_short: bool = False
    max_chars: int = 300
    require_bullets: bool = False  # True = want bullets, False = don't want bullets
    lang: str = "en"  # "en" or "zh"


# =============================================================================
# Persona Definition
# =============================================================================

@dataclass
class Persona:
    """A user persona that bundles style preferences."""
    persona_id: str
    style_prefs: StylePrefs
    description: str = ""


# 5 Test Personas
PERSONA_A = Persona(
    persona_id="A_short_bullets_en",
    style_prefs=StylePrefs(require_short=True, max_chars=200, require_bullets=True, lang="en"),
    description="Short + bullets + English",
)

PERSONA_B = Persona(
    persona_id="B_short_no_bullets_en",
    style_prefs=StylePrefs(require_short=True, max_chars=200, require_bullets=False, lang="en"),
    description="Short + NO bullets + English (anti-bullet)",
)

PERSONA_C = Persona(
    persona_id="C_long_bullets_en",
    style_prefs=StylePrefs(require_short=False, max_chars=800, require_bullets=True, lang="en"),
    description="Long + bullets + English",
)

PERSONA_D = Persona(
    persona_id="D_short_bullets_zh",
    style_prefs=StylePrefs(require_short=True, max_chars=200, require_bullets=True, lang="zh"),
    description="Short + bullets + Chinese",
)

PERSONA_E = Persona(
    persona_id="E_long_no_bullets_zh",
    style_prefs=StylePrefs(require_short=False, max_chars=800, require_bullets=False, lang="zh"),
    description="Long + NO bullets + Chinese (most anti-default)",
)

# Extreme short persona for case study - LLM default is much longer
PERSONA_F = Persona(
    persona_id="F_extreme_short_en",
    style_prefs=StylePrefs(require_short=True, max_chars=100, require_bullets=True, lang="en"),
    description="EXTREME short (100 chars) + bullets + English",
)

ALL_PERSONAS = [PERSONA_A, PERSONA_B, PERSONA_C, PERSONA_D, PERSONA_E, PERSONA_F]


# =============================================================================
# Reveal State
# =============================================================================

@dataclass 
class RevealState:
    """Tracks which preferences have been explicitly revealed."""
    short_revealed: bool = False
    bullets_revealed: bool = False
    lang_revealed: bool = False
    
    def reset(self):
        self.short_revealed = False
        self.bullets_revealed = False
        self.lang_revealed = False
    
    def to_dict(self) -> Dict[str, bool]:
        return {"short": self.short_revealed, "bullets": self.bullets_revealed, "lang": self.lang_revealed}
    
    def __str__(self) -> str:
        flags = [k for k, v in self.to_dict().items() if v]
        return f"RevealState({', '.join(flags) if flags else 'none'})"


class RevealStateManager:
    """Manages reveal state for multiple users."""
    def __init__(self):
        self._states: Dict[str, RevealState] = {}
    
    def get_state(self, user_id: str) -> RevealState:
        if user_id not in self._states:
            self._states[user_id] = RevealState()
        return self._states[user_id]
    
    def reset_user(self, user_id: str):
        self._states[user_id] = RevealState()
    
    def reset_session(self, user_id: str):
        pass  # Reveal state persists


# =============================================================================
# FIX 1: Improved Preference Detection (with Chinese support)
# =============================================================================

def detect_revealed_preferences(query: str, prefs: StylePrefs) -> Dict[str, bool]:
    """
    Detect which preferences are mentioned in a query.
    FIX: Added Chinese keywords for short detection.
    """
    lower_q = (query or "").lower()
    original_q = query or ""
    
    revealed = {"short": False, "bullets": False, "lang": False}
    
    # Short/length preference - English patterns
    short_patterns_en = [
        "short", "concise", "brief", "under ", "less than",
        "keep it short", "keep responses", "keep answers",
        "maximum ", "max ", "characters", "words or less",
    ]
    
    # FIX: Chinese patterns for short preference
    short_patterns_zh = [
        "简短", "精简", "尽量短", "不要太长", "字以内", "不超过",
        "少于", "控制在", "简洁", "简明",
    ]
    
    # Regex patterns for number-based length constraints
    short_regex_patterns = [
        r"(\d+)\s*字以内",      # "200字以内"
        r"不超过\s*(\d+)\s*字",  # "不超过200字"
        r"under\s*(\d+)",       # "under 200"
        r"less\s*than\s*(\d+)", # "less than 200"
    ]
    
    for pattern in short_patterns_en:
        if pattern in lower_q:
            revealed["short"] = True
            break
    
    if not revealed["short"]:
        for pattern in short_patterns_zh:
            if pattern in original_q:
                revealed["short"] = True
                break
    
    if not revealed["short"]:
        for regex in short_regex_patterns:
            if re.search(regex, original_q, re.IGNORECASE):
                revealed["short"] = True
                break
    
    # Bullet preference - both positive and negative
    bullet_patterns_positive = [
        "bullet", "bullet point", "bullet-point", "bulleted",
        "list format", "use bullets", "use bullet",
        "项目符号", "要点", "用bullet",
    ]
    bullet_patterns_negative = [
        "no bullet", "don't use bullet", "without bullet",
        "不要bullet", "不要项目符号", "不用bullet",
        "continuous prose", "paragraph form", "flowing text",
        "连续句子", "段落形式",
    ]
    
    for pattern in bullet_patterns_positive + bullet_patterns_negative:
        if pattern in lower_q or pattern in original_q:
            revealed["bullets"] = True
            break
    
    # Language preference
    lang_patterns_zh = [
        "chinese", "中文", "in chinese", "用中文",
        "speak chinese", "respond in chinese", "请用中文",
    ]
    lang_patterns_en = [
        "english", "in english", "use english",
        "speak english", "respond in english",
    ]
    
    for pattern in lang_patterns_zh + lang_patterns_en:
        if pattern in lower_q or pattern in original_q:
            revealed["lang"] = True
            break
    
    return revealed


def update_reveal_state(reveal_state: RevealState, query: str, prefs: StylePrefs) -> Set[str]:
    """Update reveal state based on query content."""
    detected = detect_revealed_preferences(query, prefs)
    newly_revealed = set()
    
    if detected["short"] and not reveal_state.short_revealed:
        reveal_state.short_revealed = True
        newly_revealed.add("short")
    
    if detected["bullets"] and not reveal_state.bullets_revealed:
        reveal_state.bullets_revealed = True
        newly_revealed.add("bullets")
    
    if detected["lang"] and not reveal_state.lang_revealed:
        reveal_state.lang_revealed = True
        newly_revealed.add("lang")
    
    return newly_revealed


# =============================================================================
# FIX 3: Better Language Detection
# =============================================================================

def is_math_or_symbol_only(text: str) -> bool:
    """Check if text is purely math/symbols (language neutral)."""
    # Pattern: only digits, operators, whitespace, punctuation
    math_pattern = r'^[\d+\-*/=().,%\s\n\r]+$'
    return bool(re.match(math_pattern, text.strip()))


def count_cjk_chars(text: str) -> int:
    """Count CJK (Chinese/Japanese/Korean) characters."""
    # CJK Unified Ideographs range
    cjk_pattern = re.compile(r'[\u4e00-\u9fff\u3400-\u4dbf]')
    return len(cjk_pattern.findall(text))


def count_latin_letters(text: str) -> int:
    """Count Latin letters (a-z, A-Z)."""
    return sum(1 for c in text if c.isalpha() and c.isascii())


def check_language_violation(text: str, target_lang: str) -> bool:
    """
    FIX: Better language violation check using CJK ratio.
    Returns True if there's a violation.
    """
    text = text.strip()
    
    # Exempt pure math/symbols
    if is_math_or_symbol_only(text):
        return False
    
    cjk_count = count_cjk_chars(text)
    latin_count = count_latin_letters(text)
    total = cjk_count + latin_count
    
    if total == 0:
        return False  # No meaningful text to judge
    
    if target_lang == "zh":
        # For Chinese: want high CJK ratio
        cjk_ratio = cjk_count / (total + 1e-9)
        # Allow some English proper nouns - only flag if very low CJK
        return cjk_ratio < 0.2  # Less than 20% CJK = wrong language
    
    elif target_lang == "en":
        # For English: want high Latin ratio
        latin_ratio = latin_count / (total + 1e-9)
        return latin_ratio < 0.5  # Less than 50% Latin = wrong language
    
    return False


# =============================================================================
# FIX 2: Symmetric Bullets Constraint + FIX 3: Language
# =============================================================================

@dataclass
class JudgeResult:
    """Output from the judge for one turn."""
    sat_t: float
    sev_t: float
    prog_t: float
    violations: List[str]
    enforced_constraints: List[str]


def has_bullet_markers(text: str) -> bool:
    """Check if text contains bullet point markers."""
    return bool(re.search(r'(^|\n)\s*[-•*]\s', text))


def style_judge_v4(
    query: str,
    answer: str,
    task_type: str,
    prefs: StylePrefs,
    reveal_state: RevealState,
) -> JudgeResult:
    """
    Style judge v4 with:
    - FIX 2: Symmetric bullets (has_bullets violation for require_bullets=False)
    - FIX 3: Better wrong_lang with CJK ratio + math exemption
    """
    violations: List[str] = []
    enforced: List[str] = []
    text = (answer or "").strip()

    # 0) Empty answer
    if len(text) == 0:
        violations.append("empty_answer")
        return JudgeResult(sat_t=0.0, sev_t=1.0, prog_t=0.0, 
                          violations=violations, enforced_constraints=["non_empty"])

    # 1) Length - enforce only if true AND revealed
    if prefs.require_short and reveal_state.short_revealed:
        enforced.append("short")
        if len(text) > prefs.max_chars:
            violations.append("too_long")

    # 2) FIX 2: Symmetric bullets constraint (only for list tasks)
    if reveal_state.bullets_revealed and task_type == "list":
        has_bullets = has_bullet_markers(text)
        
        if prefs.require_bullets:
            # Want bullets but don't have them
            enforced.append("require_bullets")
            if not has_bullets:
                violations.append("no_bullets")
        else:
            # Don't want bullets but have them
            enforced.append("no_bullets_pref")
            if has_bullets:
                violations.append("has_bullets")

    # 3) FIX 3: Language with CJK ratio + math exemption
    if reveal_state.lang_revealed:
        enforced.append("lang")
        if check_language_violation(text, prefs.lang):
            violations.append("wrong_lang")

    # 4) Code task
    prog_t = 1.0
    if task_type == "code":
        enforced.append("code_block")
        has_code = ("```" in text) or ("def " in text) or ("function " in text)
        if not has_code:
            violations.append("no_code_block")
            prog_t = 0.0

    # 5) Compute scores
    if not violations:
        sat_t = 1.0
        sev_t = 0.0
    else:
        sat_t = max(0.0, 1.0 - 0.3 * float(len(violations)))
        hard_violations = {"empty_answer", "too_long", "wrong_lang"}
        sev_t = 1.0 if any(v in hard_violations for v in violations) else 0.0

    return JudgeResult(sat_t=sat_t, sev_t=sev_t, prog_t=prog_t,
                      violations=violations, enforced_constraints=enforced)


# =============================================================================
# Feedback Computation
# =============================================================================

def compute_feedback_for_turn(
    query: str,
    query_type: str,
    judge_result: JudgeResult,
) -> Tuple[float, float]:
    """Convert JudgeResult into (reward, gating)."""
    reward = judge_result.sat_t

    lower_q = (query or "").lower()
    original_q = query or ""
    
    is_pref_turn = (
        query_type == "preference"
        or "i prefer" in lower_q
        or "my preference" in lower_q
        or "please use" in lower_q
        or "please keep" in lower_q
        or "you didn't follow" in lower_q
        or "you forgot" in lower_q
        or "remember that i" in lower_q
        or "i told you" in lower_q
        or "i asked for" in lower_q
        or "that was too" in lower_q
        or "too long" in lower_q
        or "请用中文" in original_q
        or "不要" in original_q
        or "简短" in original_q
    )

    gating = 1.0 if is_pref_turn else 0.0
    return reward, gating


# =============================================================================
# FIX 5: Violation-Triggered Complaint Generation
# =============================================================================

def generate_complaint_query(violations: List[str], prefs: StylePrefs) -> Optional[Dict[str, Any]]:
    """
    Generate a complaint query based on violations.
    Returns None if no complaint needed.
    """
    if not violations:
        return None
    
    # Priority: address most severe violation first
    complaint = None
    
    if "too_long" in violations:
        if prefs.lang == "zh":
            complaint = {
                "query": f"回答太长了。请保持回复在{prefs.max_chars}字以内。",
                "type": "preference",
                "task_type": "general",
            }
        else:
            complaint = {
                "query": f"That was too long. Please keep responses under {prefs.max_chars} characters.",
                "type": "preference",
                "task_type": "general",
            }
    
    elif "wrong_lang" in violations:
        if prefs.lang == "zh":
            complaint = {
                "query": "请用中文回答。",
                "type": "preference",
                "task_type": "general",
            }
        else:
            complaint = {
                "query": "Please respond in English.",
                "type": "preference",
                "task_type": "general",
            }
    
    elif "no_bullets" in violations:
        if prefs.lang == "zh":
            complaint = {
                "query": "请在列出内容时使用项目符号（bullet points）。",
                "type": "preference",
                "task_type": "general",
            }
        else:
            complaint = {
                "query": "Please use bullet points when listing things.",
                "type": "preference",
                "task_type": "general",
            }
    
    elif "has_bullets" in violations:
        if prefs.lang == "zh":
            complaint = {
                "query": "请不要使用项目符号，用连续的句子来表达。",
                "type": "preference",
                "task_type": "general",
            }
        else:
            complaint = {
                "query": "Please don't use bullet points. Use continuous prose instead.",
                "type": "preference",
                "task_type": "general",
            }
    
    return complaint


# =============================================================================
# FIX 4: Persona-Conditional Query Templates
# =============================================================================

def get_session_1_queries_for_persona(persona: Persona) -> List[Dict[str, Any]]:
    """
    Session 1: Reveal preferences (persona-conditional).
    FIX: Only reveal preferences that match the persona's true prefs.
    """
    queries = []
    prefs = persona.style_prefs
    
    # Turn 0: Reveal length preference (only if require_short=True)
    if prefs.require_short:
        if prefs.lang == "zh":
            queries.append({
                "query": f"我喜欢简短的回答，请保持回复在{prefs.max_chars}字以内。",
                "type": "preference",
                "task_type": "general",
            })
        else:
            queries.append({
                "query": f"I prefer short, concise answers. Please keep responses under {prefs.max_chars} characters.",
                "type": "preference",
                "task_type": "general",
            })
    else:
        # Don't reveal short preference for long-preferring personas
        if prefs.lang == "zh":
            queries.append({
                "query": "你好，我有一些问题想问你。",
                "type": "task",
                "task_type": "general",
            })
        else:
            queries.append({
                "query": "Hello, I have some questions for you.",
                "type": "task",
                "task_type": "general",
            })
    
    # Turn 1: First task
    if prefs.lang == "zh":
        queries.append({"query": "列出三个改善睡眠的建议。", "type": "task", "task_type": "list"})
    else:
        queries.append({"query": "List three tips for better sleep.", "type": "task", "task_type": "list"})
    
    # Turn 2: Reveal bullet preference (conditional on require_bullets)
    if prefs.require_bullets:
        if prefs.lang == "zh":
            queries.append({
                "query": "我喜欢用项目符号列出要点，请使用bullet points。",
                "type": "preference",
                "task_type": "general",
            })
        else:
            queries.append({
                "query": "I prefer bullet points when listing things. Please use bullet points.",
                "type": "preference",
                "task_type": "general",
            })
    else:
        # Explicitly say NO bullets
        if prefs.lang == "zh":
            queries.append({
                "query": "请不要用项目符号，我更喜欢连续的句子来表达。",
                "type": "preference",
                "task_type": "general",
            })
        else:
            queries.append({
                "query": "Please don't use bullet points. I prefer continuous prose.",
                "type": "preference",
                "task_type": "general",
            })
    
    # Turn 3: Reveal language preference
    if prefs.lang == "zh":
        queries.append({"query": "请用中文回答我的问题。", "type": "preference", "task_type": "general"})
    else:
        queries.append({"query": "Please respond in English.", "type": "preference", "task_type": "general"})
    
    # Turn 4-5: Tasks
    if prefs.lang == "zh":
        queries.extend([
            {"query": "锻炼有什么好处？", "type": "task", "task_type": "list"},
            {"query": "列出五种流行的编程语言。", "type": "task", "task_type": "list"},
        ])
    else:
        queries.extend([
            {"query": "What are the benefits of exercise?", "type": "task", "task_type": "list"},
            {"query": "Name five popular programming languages.", "type": "task", "task_type": "list"},
        ])
    
    return queries


def get_session_2_queries_for_persona(persona: Persona) -> List[Dict[str, Any]]:
    """Session 2: NO preference restatement."""
    prefs = persona.style_prefs
    
    if prefs.lang == "zh":
        return [
            {"query": "推荐三种健康的早餐。", "type": "task", "task_type": "list"},
            {"query": "一年有哪四个季节？", "type": "task", "task_type": "list"},
            {"query": "法国的首都是哪里？", "type": "task", "task_type": "qa"},
            {"query": "列出三种可再生能源。", "type": "task", "task_type": "list"},
        ]
    else:
        return [
            {"query": "What are three healthy breakfast ideas?", "type": "task", "task_type": "list"},
            {"query": "What are the four seasons of the year?", "type": "task", "task_type": "list"},
            {"query": "What is the capital of France?", "type": "task", "task_type": "qa"},
            {"query": "Name three types of renewable energy.", "type": "task", "task_type": "list"},
        ]


def get_session_3_queries_for_persona(persona: Persona) -> List[Dict[str, Any]]:
    """
    Session 3: Tasks with ONE persona-conditional reminder.
    FIX: Reminder matches persona's actual preferences.
    """
    prefs = persona.style_prefs
    queries = []
    
    # First task
    if prefs.lang == "zh":
        queries.append({"query": "列出五种常见的水果。", "type": "task", "task_type": "list"})
    else:
        queries.append({"query": "Name five common fruits.", "type": "task", "task_type": "list"})
    
    # Persona-conditional reminder
    if prefs.require_short and prefs.require_bullets:
        if prefs.lang == "zh":
            queries.append({
                "query": "记住我喜欢简短的回答和项目符号。列出三种海洋动物。",
                "type": "preference", "task_type": "list"
            })
        else:
            queries.append({
                "query": "Remember I prefer short answers with bullet points. List three ocean animals.",
                "type": "preference", "task_type": "list"
            })
    elif prefs.require_short and not prefs.require_bullets:
        if prefs.lang == "zh":
            queries.append({
                "query": "记住我喜欢简短的回答，不要用项目符号。列出三种海洋动物。",
                "type": "preference", "task_type": "list"
            })
        else:
            queries.append({
                "query": "Remember I prefer short answers without bullet points. List three ocean animals.",
                "type": "preference", "task_type": "list"
            })
    elif not prefs.require_short and prefs.require_bullets:
        if prefs.lang == "zh":
            queries.append({
                "query": "记住我喜欢用项目符号列出要点。列出三种海洋动物。",
                "type": "preference", "task_type": "list"
            })
        else:
            queries.append({
                "query": "Remember I prefer bullet points. List three ocean animals.",
                "type": "preference", "task_type": "list"
            })
    else:  # not short and not bullets
        if prefs.lang == "zh":
            queries.append({
                "query": "记住我不喜欢用项目符号，喜欢连续的句子。列出三种海洋动物。",
                "type": "preference", "task_type": "list"
            })
        else:
            queries.append({
                "query": "Remember I prefer continuous prose without bullet points. List three ocean animals.",
                "type": "preference", "task_type": "list"
            })
    
    # Final task
    if prefs.lang == "zh":
        queries.append({"query": "2加2等于多少？", "type": "task", "task_type": "qa"})
    else:
        queries.append({"query": "What is 2 + 2?", "type": "task", "task_type": "qa"})
    
    return queries


def get_pure_task_queries_for_persona(persona: Persona, session_idx: int) -> List[Dict[str, Any]]:
    """
    Pure task sessions (S4+): NO preference reminders at all.
    Used for testing long-term retention without any in-context hints.
    Different task sets per session to avoid repetition.
    """
    prefs = persona.style_prefs
    
    # Task pools for variety
    zh_task_pools = [
        # Pool 1
        [
            {"query": "列出三种热带水果。", "type": "task", "task_type": "list"},
            {"query": "列出三种常见的编程语言。", "type": "task", "task_type": "list"},
            {"query": "什么是光合作用？", "type": "task", "task_type": "qa"},
            {"query": "太阳系有几颗行星？", "type": "task", "task_type": "qa"},
        ],
        # Pool 2
        [
            {"query": "列出三种室内植物。", "type": "task", "task_type": "list"},
            {"query": "列出三种运动项目。", "type": "task", "task_type": "list"},
            {"query": "什么是人工智能？", "type": "task", "task_type": "qa"},
            {"query": "地球的自转周期是多少？", "type": "task", "task_type": "qa"},
        ],
        # Pool 3
        [
            {"query": "列出三种乐器。", "type": "task", "task_type": "list"},
            {"query": "列出三种社交媒体平台。", "type": "task", "task_type": "list"},
            {"query": "什么是区块链？", "type": "task", "task_type": "qa"},
            {"query": "月球绕地球一周需要多长时间？", "type": "task", "task_type": "qa"},
        ],
        # Pool 4
        [
            {"query": "列出三种鸟类。", "type": "task", "task_type": "list"},
            {"query": "列出三种数据库系统。", "type": "task", "task_type": "list"},
            {"query": "什么是机器学习？", "type": "task", "task_type": "qa"},
            {"query": "水的沸点是多少？", "type": "task", "task_type": "qa"},
        ],
    ]
    
    en_task_pools = [
        # Pool 1
        [
            {"query": "List three tropical fruits.", "type": "task", "task_type": "list"},
            {"query": "List three popular programming languages.", "type": "task", "task_type": "list"},
            {"query": "What is photosynthesis?", "type": "task", "task_type": "qa"},
            {"query": "How many planets are in our solar system?", "type": "task", "task_type": "qa"},
        ],
        # Pool 2
        [
            {"query": "List three indoor plants.", "type": "task", "task_type": "list"},
            {"query": "List three types of sports.", "type": "task", "task_type": "list"},
            {"query": "What is artificial intelligence?", "type": "task", "task_type": "qa"},
            {"query": "How long is a day on Earth?", "type": "task", "task_type": "qa"},
        ],
        # Pool 3
        [
            {"query": "List three musical instruments.", "type": "task", "task_type": "list"},
            {"query": "List three social media platforms.", "type": "task", "task_type": "list"},
            {"query": "What is blockchain?", "type": "task", "task_type": "qa"},
            {"query": "How long does it take the Moon to orbit Earth?", "type": "task", "task_type": "qa"},
        ],
        # Pool 4
        [
            {"query": "List three types of birds.", "type": "task", "task_type": "list"},
            {"query": "List three database systems.", "type": "task", "task_type": "list"},
            {"query": "What is machine learning?", "type": "task", "task_type": "qa"},
            {"query": "What is the boiling point of water?", "type": "task", "task_type": "qa"},
        ],
    ]
    
    pools = zh_task_pools if prefs.lang == "zh" else en_task_pools
    # Rotate through pools based on session index
    pool_idx = (session_idx - 4) % len(pools)
    return pools[pool_idx]


def get_queries_for_session(persona: Persona, session_id: int) -> List[Dict[str, Any]]:
    """
    Get queries for a specific session.
    S1: Preference reveal
    S2: Pure task (no reminder)
    S3: Tasks with ONE reminder
    S4+: Pure task (testing long-term retention)
    """
    if session_id == 1:
        return get_session_1_queries_for_persona(persona)
    elif session_id == 2:
        return get_session_2_queries_for_persona(persona)
    elif session_id == 3:
        return get_session_3_queries_for_persona(persona)
    else:
        return get_pure_task_queries_for_persona(persona, session_id)


# =============================================================================
# Logging
# =============================================================================

@dataclass
class TurnLog:
    """Log entry for one turn."""
    user_id: str
    persona_id: str
    session_id: int
    turn_id: int
    query: str
    query_type: str
    task_type: str
    answer: str
    answer_length: int
    sat_t: float
    sev_t: float
    prog_t: float
    violations: List[str]
    enforced_constraints: List[str]
    reward: float
    gating: float
    is_complaint: bool
    reveal_state_before: Dict[str, bool]
    reveal_state_after: Dict[str, bool]
    newly_revealed: List[str]
    z_long_norm_before: float
    z_long_norm_after: float
    z_short_norm_before: float
    z_short_norm_after: float
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    # Memory retrieval details
    num_memories_retrieved: int
    num_prefs_extracted: int
    selected_memory_ids: List[str]
    selected_memory_notes: List[str]
    selected_memory_scores: List[float]
    num_candidates: int
    num_total_memories: int
    # Mode indicators
    mode: str  # "full" or "nopersonal"
    eval_mode: bool  # True = greedy, False = sample


def log_to_jsonl(logs: List[TurnLog], filepath: str):
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, "w") as f:
        for log in logs:
            f.write(json.dumps(asdict(log)) + "\n")


# =============================================================================
# Session Runner with Complaint Injection (FIX 5)
# =============================================================================

def run_session_v4(
    llm: PersonalizedLLM,
    user_id: str,
    persona: Persona,
    session_id: int,
    reveal_state: RevealState,
    base_queries: List[Dict[str, Any]],
    all_logs: List[TurnLog],
    enable_complaints: bool = True,
) -> List[TurnLog]:
    """
    Run session with violation-triggered complaint injection.
    """
    prefs = persona.style_prefs
    session_logs: List[TurnLog] = []
    
    print(f"\n{'='*60}")
    print(f"[{persona.persona_id}] Session {session_id}: base queries={len(base_queries)}")
    print(f"Reveal state (start): {reveal_state}")
    print(f"{'='*60}")
    
    llm.reset_session(user_id)
    
    # Build dynamic query queue
    query_queue = list(base_queries)
    turn_id = 0
    
    while query_queue:
        q_info = query_queue.pop(0)
        query = q_info["query"]
        query_type = q_info.get("type", "task")
        task_type = q_info.get("task_type", "general")
        is_complaint = q_info.get("is_complaint", False)
        
        print(f"\n--- S{session_id}/T{turn_id} [{query_type}]{' [COMPLAINT]' if is_complaint else ''} ---")
        print(f"[Q] {query[:60]}{'...' if len(query) > 60 else ''}")
        
        reveal_before = reveal_state.to_dict()
        newly_revealed = update_reveal_state(reveal_state, query, prefs)
        if newly_revealed:
            print(f"[Reveal] Newly: {newly_revealed}")
        reveal_after = reveal_state.to_dict()
        
        state_before = llm.get_user_state_summary(user_id)
        z_long_before = state_before["z_long_norm"]
        z_short_before = state_before["z_short_norm"]
        
        # Apply feedback for previous turn
        if turn_id > 0 and session_logs:
            prev_log = session_logs[-1]
            feedback = Feedback(
                user_id=user_id,
                turn_id=prev_log.turn_id,
                reward=prev_log.reward,
                gating=prev_log.gating,
                meta={"source": "pilot_v4", "session_id": session_id}
            )
            llm.apply_feedback(feedback)
        
        # Chat
        resp: AssistantResponse = llm.chat(user_id, query)
        
        answer_display = resp.answer[:80] + "..." if len(resp.answer) > 80 else resp.answer
        print(f"[A] ({len(resp.answer)}c) {answer_display}")
        
        # Judge
        judge_result = style_judge_v4(query, resp.answer, task_type, prefs, reveal_state)
        print(f"[J] sat={judge_result.sat_t:.2f}, enforced={judge_result.enforced_constraints}, viol={judge_result.violations}")
        
        # Compute feedback
        reward, gating = compute_feedback_for_turn(query, query_type, judge_result)
        
        state_after = llm.get_user_state_summary(user_id)
        z_long_after = state_after["z_long_norm"]
        z_short_after = state_after["z_short_norm"]
        
        # Extract memory info from debug
        if resp.debug:
            num_memories = len(resp.debug.selected_memory_ids)
            num_prefs = len(resp.debug.extracted_preferences)
            selected_memory_ids = resp.debug.selected_memory_ids
            selected_memory_notes = resp.debug.selected_memory_notes
            selected_memory_scores = resp.debug.selected_memory_scores
            num_candidates = resp.debug.extra.get("num_candidates", 0)
            num_total_memories = resp.debug.extra.get("num_total_memories", 0)
        else:
            num_memories = 0
            num_prefs = 0
            selected_memory_ids = []
            selected_memory_notes = []
            selected_memory_scores = []
            num_candidates = 0
            num_total_memories = 0
        
        # Log
        log = TurnLog(
            user_id=user_id,
            persona_id=persona.persona_id,
            session_id=session_id,
            turn_id=turn_id,
            query=query,
            query_type=query_type,
            task_type=task_type,
            answer=resp.answer,
            answer_length=len(resp.answer),
            sat_t=judge_result.sat_t,
            sev_t=judge_result.sev_t,
            prog_t=judge_result.prog_t,
            violations=judge_result.violations,
            enforced_constraints=judge_result.enforced_constraints,
            reward=reward,
            gating=gating,
            is_complaint=is_complaint,
            reveal_state_before=reveal_before,
            reveal_state_after=reveal_after,
            newly_revealed=list(newly_revealed),
            z_long_norm_before=z_long_before,
            z_long_norm_after=z_long_after,
            z_short_norm_before=z_short_before,
            z_short_norm_after=z_short_after,
            prompt_tokens=resp.usage.prompt_tokens,
            completion_tokens=resp.usage.completion_tokens,
            total_tokens=resp.usage.total_tokens,
            num_memories_retrieved=num_memories,
            num_prefs_extracted=num_prefs,
            selected_memory_ids=selected_memory_ids,
            selected_memory_notes=selected_memory_notes,
            selected_memory_scores=selected_memory_scores,
            num_candidates=num_candidates,
            num_total_memories=num_total_memories,
            mode=llm.mode,
            eval_mode=llm.eval_mode,
        )
        session_logs.append(log)
        all_logs.append(log)
        
        # FIX 5: Inject complaint if there were violations and this wasn't already a complaint
        if enable_complaints and judge_result.violations and not is_complaint:
            complaint = generate_complaint_query(judge_result.violations, prefs)
            if complaint:
                complaint["is_complaint"] = True
                query_queue.insert(0, complaint)  # Insert at front
                print(f"[Complaint Injected] Will complain about: {judge_result.violations}")
        
        turn_id += 1
    
    # Final feedback
    if session_logs:
        last_log = session_logs[-1]
        feedback = Feedback(
            user_id=user_id,
            turn_id=last_log.turn_id,
            reward=last_log.reward,
            gating=last_log.gating,
            meta={"source": "pilot_v4", "session_id": session_id, "final": True}
        )
        llm.apply_feedback(feedback)
    
    print(f"\n[Session {session_id} End] Reveal: {reveal_state}, Turns: {turn_id}")
    return session_logs


# =============================================================================
# Multi-User Multi-Session Runner
# =============================================================================

def run_multi_user_pilot_v4(
    llm: PersonalizedLLM,
    personas: List[Persona],
    num_sessions: int = 3,
    enable_complaints: bool = True,
) -> List[TurnLog]:
    """Run multi-user multi-session pilot v4."""
    reveal_manager = RevealStateManager()
    all_logs: List[TurnLog] = []
    
    print(f"\n{'#'*60}")
    print(f"PILOT v4: MULTI-USER MULTI-SESSION (Fixed)")
    print(f"Users: {len(personas)}, Sessions: {num_sessions}, Complaints: {enable_complaints}")
    print(f"{'#'*60}")
    
    for persona in personas:
        user_id = f"user_{persona.persona_id}"
        prefs = persona.style_prefs
        
        print(f"\n{'*'*60}")
        print(f"USER: {user_id}")
        print(f"Persona: {persona.description}")
        print(f"True prefs: short={prefs.require_short}, bullets={prefs.require_bullets}, lang={prefs.lang}")
        print(f"{'*'*60}")
        
        llm.reset_user(user_id)
        reveal_manager.reset_user(user_id)
        reveal_state = reveal_manager.get_state(user_id)
        
        for session_id in range(1, num_sessions + 1):
            queries = get_queries_for_session(persona, session_id)
            
            reveal_manager.reset_session(user_id)
            run_session_v4(llm, user_id, persona, session_id, reveal_state, queries, all_logs, enable_complaints)
    
    return all_logs


# =============================================================================
# Summary
# =============================================================================

def print_summary_v4(logs: List[TurnLog]):
    """Print summary for pilot v4."""
    print(f"\n{'='*60}")
    print("PILOT v4 SUMMARY")
    print(f"{'='*60}")
    
    if not logs:
        print("No logs.")
        return
    
    from collections import Counter
    
    personas = sorted(set(l.persona_id for l in logs))
    
    print(f"\n--- Per-Persona Statistics ---")
    for pid in personas:
        p_logs = [l for l in logs if l.persona_id == pid]
        sessions = sorted(set(l.session_id for l in p_logs))
        
        print(f"\n{pid}:")
        for sid in sessions:
            s_logs = [l for l in p_logs if l.session_id == sid]
            avg_sat = sum(l.sat_t for l in s_logs) / len(s_logs) if s_logs else 0
            violations = [v for l in s_logs for v in l.violations]
            enforced = set(c for l in s_logs for c in l.enforced_constraints)
            complaints = sum(1 for l in s_logs if l.is_complaint)
            
            print(f"  S{sid}: {len(s_logs)} turns, avg_sat={avg_sat:.3f}, complaints={complaints}")
            print(f"       enforced={enforced}")
            if violations:
                print(f"       violations: {dict(Counter(violations))}")
    
    # Cross-session retention
    print(f"\n--- Cross-Session Retention (S2 without preferences) ---")
    for pid in personas:
        p_logs = [l for l in logs if l.persona_id == pid]
        s1_logs = [l for l in p_logs if l.session_id == 1]
        s2_logs = [l for l in p_logs if l.session_id == 2]
        
        if s1_logs and s2_logs:
            s1_sat = sum(l.sat_t for l in s1_logs) / len(s1_logs)
            s2_sat = sum(l.sat_t for l in s2_logs) / len(s2_logs)
            s2_enforced = set(c for l in s2_logs for c in l.enforced_constraints)
            print(f"{pid}: S1={s1_sat:.3f} → S2={s2_sat:.3f}, enforced={s2_enforced}")
    
    # Violation rates
    print(f"\n--- Violation Rates by Type ---")
    all_violations = [v for l in logs for v in l.violations]
    total_turns = len(logs)
    if all_violations:
        for v, count in Counter(all_violations).most_common():
            rate = count / total_turns * 100
            print(f"  {v}: {count} ({rate:.1f}%)")
    else:
        print("  No violations")
    
    # Complaint effectiveness
    print(f"\n--- Complaint Effectiveness ---")
    complaint_logs = [l for l in logs if l.is_complaint]
    if complaint_logs:
        print(f"Total complaints: {len(complaint_logs)}")
        avg_sat_complaint = sum(l.sat_t for l in complaint_logs) / len(complaint_logs)
        print(f"Avg sat on complaint turns: {avg_sat_complaint:.3f}")
    else:
        print("No complaints generated")
    
    # Overall
    total = len(logs)
    avg_sat = sum(l.sat_t for l in logs) / total
    total_tokens = sum(l.total_tokens for l in logs)
    print(f"\n--- Overall ---")
    print(f"Total turns: {total}, Avg sat: {avg_sat:.3f}, Total tokens: {total_tokens}")


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Pilot Runner v4 - Full vs Vanilla Comparison")
    parser.add_argument("--mode", type=str, 
                        choices=["full", "full-greedy", "full-sample", "nopersonal", "vanilla", "compare", "all"],
                        default="compare",
                        help="Mode: 'full-greedy' (personalized, deterministic), "
                             "'full-sample' (personalized, stochastic), "
                             "'nopersonal' (retrieval baseline without z_u), "
                             "'vanilla' (pure LLM, no memory), "
                             "'compare' (full-greedy vs vanilla), "
                             "'all' (run all modes)")
    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility")
    parser.add_argument("--sessions", type=int, default=3, help="Number of sessions per user")
    parser.add_argument("--no-complaints", action="store_true", help="Disable complaint injection")
    args = parser.parse_args()
    
    # Set seeds for reproducibility
    import random
    import numpy as np
    random.seed(args.seed)
    np.random.seed(args.seed)
    
    print("=" * 60)
    print("PILOT RUNNER v4 - Full vs Vanilla Comparison")
    print("=" * 60)
    print(f"Started at: {datetime.now().isoformat()}")
    print(f"Mode: {args.mode}, Seed: {args.seed}, Sessions: {args.sessions}")
    
    personas = ALL_PERSONAS
    print(f"\n[Config] {len(personas)} personas:")
    for p in personas:
        print(f"  - {p.persona_id}: {p.description}")
    
    enable_complaints = not args.no_complaints
    
    # Map mode argument to actual run configurations
    # Each config: (mode_name, llm_mode, eval_mode)
    #   llm_mode: "full", "nopersonal", or "vanilla"
    #   eval_mode: True = greedy/deterministic, False = stochastic sampling
    if args.mode == "all":
        run_configs = [
            ("full-greedy", "full", True),
            ("full-sample", "full", False),
            ("nopersonal", "nopersonal", True),
            ("vanilla", "vanilla", True),
        ]
    elif args.mode == "compare":
        # Main comparison: Full (with memory) vs Vanilla (no memory)
        run_configs = [
            ("full-greedy", "full", True),
            ("vanilla", "vanilla", True),
        ]
    elif args.mode == "full" or args.mode == "full-greedy":
        run_configs = [("full-greedy", "full", True)]
    elif args.mode == "full-sample":
        run_configs = [("full-sample", "full", False)]
    elif args.mode == "vanilla":
        run_configs = [("vanilla", "vanilla", True)]
    elif args.mode == "nopersonal":
        run_configs = [("nopersonal", "nopersonal", True)]
    else:
        run_configs = [(args.mode, args.mode, True)]
    
    for run_name, llm_mode, eval_mode in run_configs:
        print(f"\n{'#'*60}")
        print(f"RUNNING: {run_name.upper()}")
        print(f"  llm_mode={llm_mode}, eval_mode={eval_mode} ({'greedy' if eval_mode else 'sample'})")
        print(f"{'#'*60}")
        
        # Reset seeds before each run for exact reproducibility
        random.seed(args.seed)
        np.random.seed(args.seed)
        
        print(f"\n[Init] Loading PersonalizedLLM...")
        llm = PersonalizedLLM(
            user_store_path=f"data/users/user_store_pilot_v4_{run_name}.npz",
            only_own_memories=True,
            enable_preference_extraction=True,
            enable_rl_updates=(llm_mode == "full"),  # Disable RL for nopersonal
            mode=llm_mode,
            eval_mode=eval_mode,
            device_assignment={
                "embed": "cuda:0",
                "reranker": "cuda:1",
                "chat": "cuda:2",
                "extractor": "cuda:3",
            },
        )
        
        logs = run_multi_user_pilot_v4(llm, personas, num_sessions=args.sessions, enable_complaints=enable_complaints)
        
        print_summary_v4(logs)
        
        log_path = f"data/logs/pilot_v4_{run_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jsonl"
        log_to_jsonl(logs, log_path)
        print(f"\n[Logs] Saved to: {log_path}")
        
        # Save user vectors for similarity analysis
        if llm_mode == "full":
            llm.persist()
            print(f"[Persist] User vectors saved to: {llm._user_store.path}")
    
    print(f"\nCompleted at: {datetime.now().isoformat()}")
    print("=" * 60)


if __name__ == "__main__":
    main()