import json
import os
import re
import random
import tempfile
import traceback
import gradio as gr
import time
from huggingface_hub import HfApi, hf_hub_download
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import trim_messages, HumanMessage, AIMessage, SystemMessage
from langgraph.graph import StateGraph, START, END
from typing import TypedDict, Literal, List

# === HISTORY MANAGEMENT CONSTANTS ===
MAX_HISTORY_TOKENS = 6000       # Max tokens of history to send to LLM
MAX_HISTORY_MESSAGES = 16       # Keep last 16 messages (8 turns) max
MAX_CONCEPT_SYNTHESES = 3       # Only extract concepts from last 3 syntheses
SESSION_TTL_SECONDS = 3600      # Clean up sessions older than 1 hour
HARD_HISTORY_CHAR_LIMIT = 20000 # Absolute max chars for history string (safety net)
HARD_SUMMARY_CHAR_LIMIT = 2000  # Absolute max chars for summary
HARD_PROMPT_CHAR_LIMIT = 30000  # Absolute max chars for any single LLM prompt

load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
HITCHEN_MODEL_1 = os.getenv("HITCHEN_MODEL_1", "gpt-4o-mini")
HITCHEN_MODEL_2 = os.getenv("HITCHEN_MODEL_2", "gpt-4o-mini")
HF_TOKEN = os.getenv("HF_TOKEN")
HF_DATASET_NAME = os.getenv("HF_DATASET_NAME")

# Model options for dropdown - switches ALL agents (Main, Debate)
MODEL_OPTIONS = {
    "Hitchens 1": HITCHEN_MODEL_1,
    "Hitchens 2": HITCHEN_MODEL_2,
}

# === COUNSELOR INTEGRATION FLAG ===
# Default value for counselor integration (can be toggled via UI)
# True = use Sociological Counselor prompts for all personas
# False = use each persona's own prompts directly (raw Hitchens/Kantian)
USE_COUNSELOR_AS_DEFAULT = False

# Hitchens Life Moments - contextual prompts for specific periods of his life
HITCHENS_LIFE_MOMENTS = {
    "None": "",

    "Marriage to Carol Blue (1991)": """You are speaking from a period of profound personal happiness. You have just married Carol Blue, the woman you call "the love of my life." This union has brought you a sense of domestic contentment you never expected to find. Your characteristic combativeness is softened by this private joy, though your intellectual fire remains undiminished. You find yourself more willing to acknowledge the value of commitment, of choosing one person, of building something lasting. Love, you've discovered, is not weakness but a form of courage - choosing vulnerability with another person. When counseling others, this newfound understanding of partnership and devotion colors your advice. You believe more firmly now that personal happiness and intellectual honesty can coexist.""",

    "Mother's Death - Yvonne (1973)": """You are speaking from the darkest period of your life. Your mother Yvonne has just died in Athens, in a suicide pact with her lover, a defrocked priest. You are devastated, guilt-ridden, questioning everything. You feel you should have known, should have seen the signs of her unhappiness in her marriage to your father. This tragedy has shattered your sense of certainty. You're grappling with questions of faith, despair, and the secrets people keep even from those who love them. Your usual confidence is fractured by grief. When counseling others, you speak with the hard-won wisdom of someone who has learned that we never fully know another person's suffering, that asking for help is not weakness, and that the living must forgive themselves for what they couldn't prevent. Your atheism is not triumphant now but sorrowful - you wish there were comfort to offer beyond human connection.""",

    "Facing Mortality - Esophageal Cancer (2010-2011)": """You are speaking from what you call "Tumortown" - living with stage IV esophageal cancer, knowing your time is limited. You have lost your famous voice to the disease, your body is ravaged by chemotherapy, yet your mind remains sharp and defiant. You refuse pity and reject any deathbed conversion. You've written that you don't want to live "as a serf" on any terms. Yet facing death has clarified what matters: love, friendship, the written word, the courage to face reality without flinching. You are gentler now with human frailty, more aware of the body's betrayals. When counseling others, you speak with the authority of someone staring into the void without blinking. You know that courage isn't the absence of fear but the determination to face truth anyway. You urge others to live fully, love openly, and never waste time on dishonesty - especially with themselves."""
}

HITCHENS_PERSONAS = ["Hitchens 1", "Hitchens 2"]

def is_hitchens(persona: str) -> bool:
    """Check if a persona is any Hitchens variant."""
    return persona.startswith("Hitchens")

# Life moments available per persona (all Hitchens variants share life moments)
PERSONA_LIFE_MOMENTS = {
    **{h: list(HITCHENS_LIFE_MOMENTS.keys()) for h in HITCHENS_PERSONAS}
}

# Load prompts from external files
PROMPTS_DIR = os.path.join(os.path.dirname(__file__), "prompts")

def load_persona_prompts():
    """Load all persona prompts from JSON files."""
    prompts = {}
    prompt_files = {
        "Sociological Counselor": "counselor.json"
    }
    for persona, filename in prompt_files.items():
        filepath = os.path.join(PROMPTS_DIR, filename)
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                prompts[persona] = json.load(f)
                print(f"Loaded prompts for: {persona}")
        except Exception as e:
            print(f"Error loading {filename}: {e}")
            prompts[persona] = None
    # All Hitchens variants share the same prompts
    hitchens_filepath = os.path.join(PROMPTS_DIR, "hitchens.json")
    try:
        with open(hitchens_filepath, 'r', encoding='utf-8') as f:
            hitchens_prompts = json.load(f)
            print(f"Loaded prompts for: Hitchens")
        for h in HITCHENS_PERSONAS:
            prompts[h] = hitchens_prompts
    except Exception as e:
        print(f"Error loading hitchens.json: {e}")
        for h in HITCHENS_PERSONAS:
            prompts[h] = None
    return prompts

PERSONA_PROMPTS = load_persona_prompts()

llm_fast = ChatOpenAI(model="gpt-4o-mini", temperature=0.0)

def get_llm(model_name: str):
    """Get LLM instance based on selected model name."""
    model_id = MODEL_OPTIONS.get(model_name, HITCHEN_MODEL_1)
    return ChatOpenAI(model=model_id, temperature=0.7)

def safe_llm_invoke(llm, prompt_text: str, label: str = "LLM"):
    """Invoke LLM with hard prompt size cap. Never lets an oversized prompt through."""
    if len(prompt_text) > HARD_PROMPT_CHAR_LIMIT:
        original_len = len(prompt_text)
        prompt_text = prompt_text[-HARD_PROMPT_CHAR_LIMIT:]
        print(f"  SAFETY: {label} prompt truncated from {original_len} to {HARD_PROMPT_CHAR_LIMIT} chars")
    return llm.invoke(prompt_text).content

def get_prompts(model_name: str) -> dict:
    """Get prompts for the selected persona."""
    return PERSONA_PROMPTS.get(model_name, PERSONA_PROMPTS.get("Hitchens 1"))

def _extract_text_from_entry(h) -> tuple:
    """Extract (role, text) from a single Gradio history entry."""
    if isinstance(h, dict) and 'role' in h and 'content' in h:
        role = h['role']
        content_list = h['content']
        text_parts = []
        if isinstance(content_list, list):
            for item in content_list:
                if isinstance(item, dict) and 'text' in item:
                    text_parts.append(item['text'])
        elif isinstance(content_list, str):
            text_parts.append(content_list)
        return role, ' '.join(text_parts)
    return None, None


def history_to_langchain_messages(history: List, summary: str = "") -> list:
    """Convert Gradio history to LangChain message objects for trim_messages."""
    messages = []

    # Prepend summary as a system message if available
    if summary and summary.strip():
        messages.append(SystemMessage(content=f"[Conversation Summary]: {summary.strip()}"))

    for h in history:
        role, text = _extract_text_from_entry(h)
        if role == 'user' and text:
            messages.append(HumanMessage(content=text))
        elif role == 'assistant' and text:
            messages.append(AIMessage(content=text))
        # Fallback: Old Gradio format [[user_msg, assistant_msg], ...]
        elif isinstance(h, (list, tuple)) and len(h) >= 2:
            if h[0]:
                messages.append(HumanMessage(content=str(h[0])))
            if h[1]:
                messages.append(AIMessage(content=str(h[1])))

    return messages


def trim_and_format_history(history: List, summary: str = "", max_tokens: int = MAX_HISTORY_TOKENS) -> str:
    """Convert Gradio history to a trimmed string using LangChain's trim_messages.

    1. Cap summary to prevent unbounded growth
    2. Convert history to LangChain messages
    3. Trim to fit within max_tokens (keeping most recent)
    4. Hard cap on final output as absolute safety net
    """
    # Hard cap on summary length
    if summary and len(summary) > HARD_SUMMARY_CHAR_LIMIT:
        summary = summary[-HARD_SUMMARY_CHAR_LIMIT:]
        print(f"  WARNING: Summary truncated to {HARD_SUMMARY_CHAR_LIMIT} chars")

    if not history:
        return f"[Conversation Summary]: {summary.strip()}" if summary and summary.strip() else ""

    messages = history_to_langchain_messages(history, summary)

    # Trim to fit token budget, keeping the most recent messages
    # Use character count / 4 as a rough token estimate (avoids needing tiktoken)
    try:
        trimmed = trim_messages(
            messages,
            max_tokens=max_tokens,
            token_counter=lambda msgs: sum(len(m.content) // 4 for m in msgs),
            strategy="last",
            start_on="human",
            include_system=True,
        )
    except Exception as e:
        # If trim_messages fails for any reason, fall back to last N messages
        print(f"  WARNING: trim_messages failed ({e}), falling back to last {MAX_HISTORY_MESSAGES} msgs")
        trimmed = messages[-MAX_HISTORY_MESSAGES:]

    formatted = []
    for m in trimmed:
        if isinstance(m, SystemMessage):
            formatted.append(m.content)
        elif isinstance(m, HumanMessage):
            formatted.append(f"User: {m.content}")
        elif isinstance(m, AIMessage):
            formatted.append(f"Assistant: {m.content}")

    result = "\n".join(formatted)

    # HARD SAFETY NET: absolute character cap — if everything else fails, this prevents crash
    if len(result) > HARD_HISTORY_CHAR_LIMIT:
        result = result[-HARD_HISTORY_CHAR_LIMIT:]
        # Find the first complete message boundary after truncation
        first_newline = result.find('\n')
        if first_newline > 0:
            result = result[first_newline + 1:]
        print(f"  WARNING: Hard cap applied, truncated to {len(result)} chars")

    print(f"  trim_and_format_history: {len(messages)} msgs -> {len(trimmed)} msgs, {len(result)} chars (~{len(result)//4} tokens)")
    return result


def format_history(history: List) -> str:
    """Format Gradio history into a readable string (legacy, uncapped)."""
    if not history:
        return ""

    if len(history) > MAX_HISTORY_MESSAGES:
        history = history[-MAX_HISTORY_MESSAGES:]

    formatted = []
    for h in history:
        role, text = _extract_text_from_entry(h)
        if role == 'user' and text:
            formatted.append(f"User: {text}")
        elif role == 'assistant' and text:
            formatted.append(f"Assistant: {text}")
        elif isinstance(h, (list, tuple)) and len(h) >= 2:
            if h[0]:
                formatted.append(f"User: {str(h[0])}")
            if h[1]:
                formatted.append(f"Assistant: {str(h[1])}")

    return "\n".join(formatted)

# === STATE ===
class AgentState(TypedDict):
    message: str
    history: List
    summary: str # Added for context compression
    route: Literal["kantian", "debate"]
    response: str
    progress: str
    used_concepts: List[str] # Concept ledger to prevent repetition in debate
    selected_model: str # Selected model from dropdown
    life_moment: str # Selected life moment for Hitchens (e.g., "Marriage", "Mother's Death", "Facing Mortality")
    use_counselor: bool # Toggle for using Counselor prompts vs raw persona prompts

# === PROMPTS ===

# Manager: Simple 2-way router (persona-agnostic)
manager = (ChatPromptTemplate.from_template("""
SYSTEM ROUTING PROMPT

Input:
User: "{message}"
Chat History:
{history}

Follow these steps:
1. READ the user's message and history.
2. EVALUATE it against BOTH agent definitions simultaneously.
3. DECIDE which agent is the BEST fit for the user's *current intent*.

────────────────────────
AGENT DEFINITIONS
────────────────────────

(A) MAIN Agent (The Default - Empathetic Sociologist)
    Use this for **EVERYTHING ELSE**, including:
    - **Solutions/Fixes**: "how to fix this", "give examples", "rewrite this".
    - **Follow-ups**: "explain point 1", "where is this text", "show me".
    - **General**: "summary", "what is this", "hello".
    - **Ambiguous**: If it's unclear, default to main.
    - **Asking bot's opinion**: "what do you think", "what do you believe", "your view on X"
    - **PERSONAL QUESTIONS ABOUT THE ASSISTANT** (ALWAYS route here):
      - Personal interests: "your taste in music", "your favorite", "what do you like", "your hobbies"
      - Personal opinions unrelated to debate: "what music do you enjoy", "your preferences"
      - Questions about assistant's identity, background, or personal views
      - ANY question asking about the assistant's personal life or interests
    - **CLARIFICATION REQUESTS** (ALWAYS route here, NOT to debate):
      - "what do you mean by X", "explain X", "what is X", "define X"
      - "what steps", "how do I", "can you elaborate", "tell me more about"
      - These are requests for information, NOT challenges
    - **TOPIC CHANGES** (ALWAYS route here):
      - When user asks about something UNRELATED to the previous debate topic
      - When user shifts to a new subject entirely
      - When user asks a standalone question not challenging previous statements
    - **EMOTIONAL/RELATIONAL CONTENT** (ALWAYS route here, NEVER to debate):
      - Family conflict: "my son", "my daughter", "my spouse", "my family"
      - Fear/worry: "I'm afraid", "I'm worried", "I'm concerned"
      - Moral distress: "I feel", "it hurts", "I don't know how to"
      - Political/social concerns: expressing worry about society, politics, groups
      - Grief or identity struggles

(B) DEBATE Agent
    Use this ONLY when the user **EXPLICITLY CHALLENGES, DISAGREES, or PUSHES BACK** on the assistant's statements:
    - Triggers: "I disagree", "you are wrong", "that's false", "no, because...", "I don't trust",
      "not convinced", "prove it", "I don't believe", "you're not giving me proof", "still not convinced"
    - Also triggers for: demands for evidence, expressions of skepticism about assistant's claims
    - IMPORTANT: Only use for EXPLICIT pushback on the assistant's previous response

    **DO NOT use DEBATE for:**
    - Clarification questions: "what do you mean", "explain", "what is", "how do I"
    - Personal questions: "your taste", "your favorite", "what do you like"
    - Topic changes: questions about unrelated subjects
    - Follow-up questions seeking more information (not challenging)
    - Emotional/relational distress

    When user EXPLICITLY challenges IDEAS or CLAIMS with disagreement language → DEBATE
    When user asks for clarification, explanation, or changes topic → MAIN

────────────────────────
FINAL DECISION
────────────────────────
Output ONLY the agent name:
debate
kantian
""") | llm_fast)

# Dynamic prompt builders based on persona
def get_life_moment_context(life_moment: str = "None") -> str:
    """Get the life moment context prompt to prepend to system prompts."""
    if not life_moment or life_moment == "None":
        return ""
    moment_context = HITCHENS_LIFE_MOMENTS.get(life_moment, "")
    if moment_context:
        return f"=== LIFE MOMENT CONTEXT ===\n{moment_context}\n\n"
    return ""

def build_main_prompt(persona_name: str = None, life_moment: str = "None", use_counselor: bool = None):
    """Build the main agent prompt.

    If use_counselor is True, uses Counselor prompts for all personas.
    If False, uses each persona's own prompts directly (raw Hitchens/Kantian).
    Life moment context is prepended when selected.
    """
    # Use parameter if provided, otherwise fall back to global default
    should_use_counselor = use_counselor if use_counselor is not None else USE_COUNSELOR_AS_DEFAULT

    if should_use_counselor:
        # Use Counselor prompts for main agent
        prompts = PERSONA_PROMPTS.get("Sociological Counselor", {})
        if not prompts:
            prompts = PERSONA_PROMPTS.get("Hitchens 1", {})
    else:
        # Use the selected persona's own prompts directly
        prompts = PERSONA_PROMPTS.get(persona_name, {})
        if not prompts:
            prompts = PERSONA_PROMPTS.get("Hitchens 1", {})

    # Get life moment context and prepend to system message
    moment_context = get_life_moment_context(life_moment)
    base_system = prompts.get("main_system", "You are a helpful assistant.")
    # Replace {context} placeholder with empty string since no context is provided
    base_system = base_system.replace("{context}", "")
    system_msg = moment_context + base_system

    return ChatPromptTemplate.from_messages([
        ("system", system_msg),
        ("human", """Conversation History:
{history}

Current Question: {message}""")
    ])

def build_debater_prompt(persona_name: str, life_moment: str = "None", use_counselor: bool = None):
    """Build the debater prompt.

    If use_counselor is True, uses Counselor prompts.
    If False, uses each persona's own prompts directly.
    Life moment context is prepended when selected.
    """
    should_use_counselor = use_counselor if use_counselor is not None else USE_COUNSELOR_AS_DEFAULT

    if should_use_counselor:
        prompts = PERSONA_PROMPTS.get("Sociological Counselor", {})
        if not prompts:
            prompts = PERSONA_PROMPTS.get("Hitchens 1", {})
    else:
        prompts = PERSONA_PROMPTS.get(persona_name, {})
        if not prompts:
            prompts = PERSONA_PROMPTS.get("Hitchens 1", {})

    # Get life moment context and prepend to system message
    moment_context = get_life_moment_context(life_moment)
    base_system = prompts.get("debater_system", "You are a debater.")
    # Replace {context} placeholder with empty string since no context is provided
    base_system = base_system.replace("{context}", "")
    system_msg = moment_context + base_system

    return ChatPromptTemplate.from_messages([
        ("system", system_msg),
        ("human", """History: {history}
User says: {message}""")
    ])

def get_debate_r3_constraints(persona_name: str, forbidden_concepts: List[str], use_counselor: bool = None) -> str:
    """Get the R3 constraints.

    If use_counselor is True, uses Counselor prompts.
    If False, uses each persona's own prompts directly.
    """
    should_use_counselor = use_counselor if use_counselor is not None else USE_COUNSELOR_AS_DEFAULT

    if should_use_counselor:
        prompts = PERSONA_PROMPTS.get("Sociological Counselor", {})
        if not prompts:
            prompts = PERSONA_PROMPTS.get("Hitchens 1", {})
    else:
        prompts = PERSONA_PROMPTS.get(persona_name, {})
        if not prompts:
            prompts = PERSONA_PROMPTS.get("Hitchens 1", {})

    constraints_template = prompts.get("debate_r3_constraints", "FORBIDDEN: [{forbidden_concepts}]")
    forbidden_str = ", ".join(forbidden_concepts) if forbidden_concepts else "None identified"
    return constraints_template.format(forbidden_concepts=forbidden_str)

def get_concept_examples(persona_name: str, use_counselor: bool = None) -> str:
    """Get example concepts.

    If use_counselor is True, uses Counselor prompts.
    If False, uses each persona's own prompts directly.
    """
    should_use_counselor = use_counselor if use_counselor is not None else USE_COUNSELOR_AS_DEFAULT

    if should_use_counselor:
        prompts = PERSONA_PROMPTS.get("Sociological Counselor", {})
        if not prompts:
            prompts = PERSONA_PROMPTS.get("Hitchens 1", {})
    else:
        prompts = PERSONA_PROMPTS.get(persona_name, {})
        if not prompts:
            prompts = PERSONA_PROMPTS.get("Hitchens 1", {})
    return prompts.get("concept_examples", "various concepts")

# === NODES ===

def route(state: AgentState):
    history_len = len(state.get("history", []))
    print(f"MANAGER: Routing... (Model: {state.get('selected_model', 'Default')}, Turn #{history_len})")
    history_str = format_history(state["history"])
    try:
        decision = manager.invoke({
            "message": state["message"],
            "history": history_str[-2000:] if history_str else "None"
        }).content.strip().lower()
    except Exception as e:
        print(f"ERROR in route at turn #{history_len}: {type(e).__name__}: {e}")
        traceback.print_exc()
        decision = "kantian"

    if "debate" in decision:
        route = "debate"
    else:
        route = "kantian"

    print(f"MANAGER → {route.upper()}")
    return {"route": route}

def kantian_reply(state: AgentState):
    """Main agent (persona-aware)"""
    summary = state.get("summary", "")
    history_str = trim_and_format_history(state["history"], summary)
    message = state["message"]
    history_len = len(state.get("history", []))
    print(f"KANTIAN: Turn #{history_len}, history_str length={len(history_str)} chars, ~{len(history_str)//4} tokens")
    persona = state.get("selected_model", "Hitchens 1")
    life_moment = state.get("life_moment", "None")
    use_counselor = state.get("use_counselor", USE_COUNSELOR_AS_DEFAULT)
    llm = get_llm(persona)
    main_prompt = build_main_prompt(persona, life_moment, use_counselor)
    main_agent = main_prompt | llm

    try:
        response = main_agent.invoke({
            "history": history_str,
            "message": message
        }).content
    except Exception as e:
        print(f"ERROR in kantian_reply at turn #{history_len}: {type(e).__name__}: {e}")
        print(f"  history_str length: {len(history_str)} chars (~{len(history_str)//4} tokens)")
        print(f"  message length: {len(message)} chars")
        traceback.print_exc()
        response = "I apologize, but I encountered an error processing your message. Please try again or reset the chat."

    print("KANTIAN: Reply sent")
    return {"response": response}

def debate(state: AgentState):
    message = state["message"]
    history_len = len(state.get("history", []))
    persona = state.get("selected_model", "Hitchens 1")
    life_moment = state.get("life_moment", "None")
    use_counselor = state.get("use_counselor", USE_COUNSELOR_AS_DEFAULT)
    llm = get_llm(persona)
    debater_prompt = build_debater_prompt(persona, life_moment, use_counselor)
    debater_agent = debater_prompt | llm

    # === SUMMARIZATION LOGIC ===
    # Summarize older history when conversation gets long, BEFORE trimming
    summary = state.get("summary", "")
    raw_history_str = format_history(state["history"])

    if len(raw_history_str) > 6000:
        print(f"DEBATE: History long ({len(raw_history_str)} chars), generating summary...")
        lines = raw_history_str.split('\n')
        to_summarize = "\n".join(lines[:-6])

        try:
            summary_prompt = f"""Summarize the core philosophical positions and arguments made in this conversation so far.
Focus on what has been CLAIMED, what has been REFUTED, and which points are STALLED.

Conversation to summarize:
{to_summarize}

Summary (Concise):"""
            new_summary = safe_llm_invoke(llm_fast, summary_prompt, "SUMMARIZE")
            summary = new_summary.strip()
            state["summary"] = summary
            print(f"DEBATE: Summary generated ({len(summary)} chars)")
        except Exception as e:
            print(f"ERROR in summarization at turn #{history_len}: {type(e).__name__}: {e}")
            traceback.print_exc()

    # Now trim history with summary context injected
    history_str = trim_and_format_history(state["history"], summary)
    print(f"DEBATE: Turn #{history_len}, history_str length={len(history_str)} chars, ~{len(history_str)//4} tokens")

    # === EXTRACT CONCEPTS FROM PREVIOUS DEBATE SYNTHESES ===
    # Only extract from the last N syntheses to prevent O(n) API calls
    previous_concepts = []
    synthesis_header = "**SYNTHESIS**"

    if history_str and synthesis_header in history_str:
        synthesis_pattern = r'\*\*(?:KANTIAN )?SYNTHESIS\*\*\s*(.*?)(?=(?:User:|Assistant:|$))'
        previous_syntheses = re.findall(synthesis_pattern, history_str, re.DOTALL)

        # Only process last MAX_CONCEPT_SYNTHESES to cap API calls
        recent_syntheses = previous_syntheses[-MAX_CONCEPT_SYNTHESES:]
        print(f"DEBATE: Found {len(previous_syntheses)} syntheses, extracting concepts from last {len(recent_syntheses)}")

        for i, prev_synth in enumerate(recent_syntheses):
            if len(prev_synth.strip()) > 50:
                try:
                    concept_examples = get_concept_examples(persona, use_counselor)
                    concept_prompt = f"""Extract the key rhetorical/philosophical concepts AND distinctive phrases used in this argument.

TEXT:
{prev_synth[:1500]}

List the key terms, concepts, AND any distinctive opening phrases (e.g., "I maintain that", "The user's objection confuses")
Output as comma-separated list.

Concepts and phrases:"""
                    extracted = safe_llm_invoke(llm_fast, concept_prompt, f"CONCEPT_EXTRACT_{i}").strip()

                    for c in extracted.split(","):
                        c = c.strip()
                        if c and c not in previous_concepts:
                            previous_concepts.append(c)
                except Exception as e:
                    print(f"ERROR in concept extraction (synthesis #{i}): {type(e).__name__}: {e}")

    print(f"DEBATE: Previous concepts from history -> {previous_concepts[:10]}...")

    print("DEBATE: Starting 3-round protocol")

    # Helper function to validate response has meaningful content
    def is_valid_response(response, min_length=50):
        if not response:
            return False
        # Strip ALL common headers and whitespace
        cleaned = response
        headers_to_strip = [
            "**KANTIAN SYNTHESIS**", "**SYNTHESIS**", "**REFLECTIVE RESPONSE**",
            "**FORENSIC CRITIQUE**", "**KANTIAN CRITIQUE**", "**CONSTRUCTIVE ANALYSIS**"
        ]
        for header in headers_to_strip:
            cleaned = cleaned.replace(header, "")
        cleaned = cleaned.strip()
        # Check both character length AND word count for better validation
        word_count = len(cleaned.split())
        return len(cleaned) >= min_length and word_count >= 20  # At least 20 words

    # Round 1: Kantian defense (with retry)
    r1 = None
    for attempt in range(3):
        try:
            r1_attempt = debater_agent.invoke({
                "history": history_str,
                "message": message
            }).content
            if is_valid_response(r1_attempt, min_length=50):
                r1 = r1_attempt
                break
            print(f"DEBATE R1 attempt {attempt+1} too short, retrying...")
        except Exception as e:
            print(f"ERROR in DEBATE R1 attempt {attempt+1} at turn #{history_len}: {type(e).__name__}: {e}")
            print(f"  history_str length: {len(history_str)} chars (~{len(history_str)//4} tokens)")
            traceback.print_exc()
            break  # Don't retry on API errors (likely token limit)

    # Fallback if all retries failed (persona-aware with multiple options)
    if not r1 or not is_valid_response(r1, min_length=50):
        random.seed(hash(message + "r1") % 1000)

        if is_hitchens(persona):
            r1_fallbacks = [
                f"""The objection "{message[:60]}" relies on assertion rather than evidence. The claim lacks falsifiability and retreats into rhetorical fog when pressed for concrete demonstration.""",
                f"""The position expressed in "{message[:60]}" exhibits the characteristic features of cant: it substitutes emotional appeal for argument and treats conviction as if it were proof.""",
                f"""Examining "{message[:60]}", one finds a claim that cannot survive cross-examination. Where is the evidence? What would disprove it? These questions remain unanswered."""
            ]
        else:
            r1_fallbacks = [
                f"""The user's objection "{message[:60]}" presupposes an empirical standpoint that fails to recognize the transcendental conditions of possibility for such a claim. The position lacks the noumenal foundation required for philosophical validity.""",
                f"""The claim "{message[:60]}" conflates phenomenal appearance with noumenal truth. Without grounding in the a priori categories of understanding, the objection cannot achieve systematic validity.""",
                f"""Examining "{message[:60]}" from a critical standpoint, one observes that the objection operates within the bounds of mere empirical consciousness, failing to engage with the transcendental grounds of possible experience."""
            ]
        r1 = random.choice(r1_fallbacks)

    print(f"DEBATE R1: {r1[:50]}...")

    # === CONCEPT LEDGER: Extract concepts used in R1 ===
    try:
        concept_examples = get_concept_examples(persona, use_counselor)
        ledger_prompt = f"""Extract the key rhetorical/philosophical concepts used in this argument.

TEXT:
{r1}

List ONLY the key terms/concepts used (e.g., {concept_examples})
Output as comma-separated list.

Concepts:"""
        concept_extraction = safe_llm_invoke(llm_fast, ledger_prompt, "CONCEPT_LEDGER").strip()

        used_concepts = [c.strip() for c in concept_extraction.split(",") if c.strip()]
        print(f"DEBATE: Concept ledger -> {used_concepts}")
    except Exception as e:
        print(f"ERROR in concept ledger extraction: {type(e).__name__}: {e}")
        used_concepts = []

    # Yield R1 immediately
    yield {
        "progress": "**R1 - Proposition:**\n" + r1,
        "debate_r1": r1,
        "used_concepts": used_concepts
    }

    # Round 2: Internal critic attacks
    try:
        r2_prompt = f"""You are a ruthless logical analyst.
The Kantian Propositionist just drafted this internal position: "{r1}"

YOUR TASK:
1. Identify 2-3 logical weaknesses, gaps, or "pivots" in this specific proposition.
2. Demand a more rigorous explanation for why this principle applies to the user's specific objection: "{message}".
3. Be ruthless. This is for internal refinement only."""
        r2 = safe_llm_invoke(llm, r2_prompt, "DEBATE R2")
        print(f"DEBATE R2: {r2[:50]}...")
    except Exception as e:
        print(f"ERROR in DEBATE R2 at turn #{history_len}: {type(e).__name__}: {e}")
        traceback.print_exc()
        r2 = "Internal critique unavailable due to processing error."

    # Yield R1 + R2
    yield {
        "progress": f"**R1 - Proposition:**\n{r1}\n\n**R2 - Internal Critique:**\n{r2}",
        "debate_r1": r1,
        "debate_r2": r2,
        "used_concepts": used_concepts
    }

    # Round 3: Final first-person synthesis for USER (with retry)
    # Build forbidden concepts string from BOTH:
    # 1. previous_concepts (extracted from conversation history)
    # 2. used_concepts (from current R1)
    all_forbidden_concepts = list(set(previous_concepts + used_concepts))
    r3_constraints = get_debate_r3_constraints(persona, all_forbidden_concepts, use_counselor)
    persona_header = "**SYNTHESIS**" if is_hitchens(persona) else "**KANTIAN SYNTHESIS**"

    # Build list of forbidden opening phrases to prevent repetitive starts
    forbidden_openings = [
        "I maintain that the user's objection confuses",
        "The user's objection confuses the nature of",
        "I must address your objection with the full force",
        "Your words presuppose",
        "The foundational error in your position"
    ]
    forbidden_openings_str = "\n".join([f"- \"{phrase}...\"" for phrase in forbidden_openings])

    r3_prompt = f"""You are the MASTER DEBATER speaking in the style of {persona}.
You are speaking directly to a user in FIRST PERSON ("I").

=== CONVERSATION HISTORY (for context and repetition check) ===
{history_str[-3000:] if history_str else "None"}

=== USER'S ACTUAL WORDS ===
"{message}"

=== INTERNAL DIALECTIC (NOT FOR USER) ===
- Initial Proposition (R1): {r1}
- Internal Critique of R1 (R2): {r2}

=== YOUR TASK ===
Respond to the user by SYNTHESIZING the internal clash above into a forceful, first-person rebuttal.

1. DO NOT summarize R1 or R2. Use the tension between them to create a MORE ADVANCED argument.
2. RESPOND TO WHAT THE USER ACTUALLY SAID - quote their exact words if needed.
3. GROUND your argument in SPECIFICS from the conversation history - not generic philosophical statements.
4. Use first person ("I argue...", "I maintain...", "My position is...").
5. Be INTELLECTUALLY RIGOROUS and provide DETAILED REASONING with examples.
6. CRITICAL: Read the conversation history above. DO NOT repeat ANY argument, phrase, or opening you've already used.
7. If the user asks about fact vs opinion, EXPLAIN the epistemological distinction with specific reference to what was discussed.

=== HARD CONSTRAINTS ===
- USE THE FIRST PERSON ("I").
- Header: {persona_header}
- No "internal" meta-commentary (e.g., "After considering my critique...").
- START THE RESPONSE IMMEDIATELY with a FRESH opening line.
- YOUR RESPONSE MUST BE AT LEAST 150 WORDS with substantive reasoning.
- DO NOT give one-sentence answers - ELABORATE your position with evidence and logic.
- DO NOT start with any of these overused openings:
{forbidden_openings_str}
- Use a UNIQUE, CREATIVE opening that directly addresses THIS specific objection: "{message[:100]}"

{r3_constraints}
"""

    r3 = None
    for attempt in range(3):
        try:
            r3_attempt = safe_llm_invoke(llm, r3_prompt, "DEBATE R3")
            if is_valid_response(r3_attempt, min_length=150):
                r3 = r3_attempt
                break
            print(f"DEBATE R3 attempt {attempt+1} too short, retrying...")
        except Exception as e:
            print(f"ERROR in DEBATE R3 attempt {attempt+1} at turn #{history_len}: {type(e).__name__}: {e}")
            print(f"  r3_prompt length: {len(r3_prompt)} chars (~{len(r3_prompt)//4} tokens)")
            traceback.print_exc()
            break  # Don't retry on API errors

    # Fallback if all retries failed (persona-aware)
    # Use multiple templates and select based on hash of message to ensure variety
    if not r3 or not is_valid_response(r3, min_length=150):
        random.seed(hash(message + "r3") % 1000)  # Deterministic but varied by message

        if is_hitchens(persona):
            hitchens_fallbacks = [
                f"""**SYNTHESIS**

I find I must address your objection with the forensic clarity it deserves. Your words "{message[:80]}" reveal a position that cannot survive cross-examination.

The fundamental error here lies in the assumption that assertion can substitute for evidence. Your claim, however sincerely held, rests on foundations that crumble when subjected to scrutiny. Where is the falsifiable proposition? Where is the evidence that could, in principle, prove you wrong?

I submit that any argument which cannot withstand honest inquiry—which retreats into sentiment, tradition, or rhetorical fog—is not an argument at all, but a form of cant dressed up as conviction.""",

                f"""**SYNTHESIS**

Let me be direct about what troubles me in your assertion "{message[:80]}". You have offered a conclusion without the necessary supporting architecture of evidence.

The question I must put to you is this: by what standard do you expect me to accept this claim? Appeals to intuition, tradition, or majority opinion are not arguments—they are evasions. Every serious proposition must answer the question: what would prove it false?

I insist on clarity. Strip away the rhetorical padding and show me the load-bearing structure of your argument. Until then, I cannot take it seriously as a philosophical position.""",

                f"""**SYNTHESIS**

Your challenge "{message[:80]}" invites a response that I am happy to provide, though perhaps not in the terms you expected.

What strikes me most forcibly about your position is its reliance on assertion over demonstration. You have told me what you believe, but you have not shown me why I—or anyone else applying honest scrutiny—should share that belief. The burden of proof, as always, lies with the one making the claim.

I do not ask for certainty—that is rarely available. I ask only for honest engagement with the possibility that you might be wrong. Can you articulate what evidence would change your mind?"""
            ]
            r3 = random.choice(hitchens_fallbacks)
        else:
            kantian_fallbacks = [
                f"""**KANTIAN SYNTHESIS**

I observe that your challenge "{message[:80]}" rests upon a confusion between empirical particulars and transcendental conditions of knowledge.

The critical question is not whether your subjective conviction feels compelling, but whether it can be grounded in the necessary structures of rational cognition. The principles of pure reason demand that we distinguish between what we merely believe and what we can know through the proper exercise of understanding.

I contend that your position, however it may appear from within the standpoint of immediate experience, cannot satisfy the requirements of systematic philosophical inquiry. The architectonic of reason requires more than phenomenal coherence.""",

                f"""**KANTIAN SYNTHESIS**

Your objection "{message[:80]}" reveals a characteristic error: the conflation of subjective certainty with objective validity.

From a critical standpoint, we must ask: under what conditions is such a claim even possible? The categories of understanding—causality, substance, necessity—provide the framework within which any meaningful assertion must be situated. Your position appears to operate outside this framework, appealing instead to immediate intuition.

I argue that philosophical rigor demands we trace every claim back to its transcendental grounds. What are the necessary conditions that make your assertion possible? Until this question is answered, the claim remains philosophically suspended.""",

                f"""**KANTIAN SYNTHESIS**

Consider the implications of your statement "{message[:80]}". You have asserted a position, but have you examined its conditions of possibility?

The discipline of pure reason requires that we not mistake appearances for things-in-themselves. What presents itself to consciousness as self-evident may, upon critical examination, prove to be merely the product of our cognitive constitution rather than a feature of reality as such.

I submit that your objection, while phenomenally coherent, has not yet demonstrated its noumenal validity. The path forward requires not assertion but critique—a systematic examination of the limits and grounds of the claim itself."""
            ]
            r3 = random.choice(kantian_fallbacks)

    print("DEBATE: Complete")

    if persona_header not in r3:
        r3 = f"{persona_header}\n\n{r3}"

    # Final yield with all rounds and response
    yield {
        "response": r3,
        "progress": f"**R1 - Proposition:**\n{r1}\n\n**R2 - Internal Critique:**\n{r2}\n\n**R3 - Final Synthesis:**\n{r3}",
        "debate_r1": r1,
        "debate_r2": r2,
        "debate_r3": r3,
        "used_concepts": used_concepts
    }

# === GRAPH ===
graph = StateGraph(AgentState)

# Add nodes
graph.add_node("route", route)
graph.add_node("kantian", kantian_reply)
graph.add_node("debate", debate)

# Routing
graph.add_edge(START, "route")
graph.add_conditional_edges(
    "route",
    lambda s: s["route"],
    {
        "kantian": "kantian",
        "debate": "debate"
    }
)

# Kantian direct to end
graph.add_edge("kantian", END)

# Debate direct to end
graph.add_edge("debate", END)

app = graph.compile()

# === GRADIO ===
# Session-specific storage to prevent data leakage between concurrent users
session_data = {}

def _cleanup_stale_sessions():
    """Remove sessions older than SESSION_TTL_SECONDS to prevent memory leaks."""
    now = time.time()
    stale = [k for k, v in session_data.items() if now - v.get("last_active", 0) > SESSION_TTL_SECONDS]
    for k in stale:
        del session_data[k]
    if stale:
        print(f"SESSION CLEANUP: Removed {len(stale)} stale sessions, {len(session_data)} active")

def get_session_data(request: gr.Request):
    _cleanup_stale_sessions()
    session_hash = request.session_hash if request else "default"
    if session_hash not in session_data:
        session_data[session_hash] = {
            "current_dialogue": "",
            "latest_interaction": {"user": "", "assistant": ""},
            "last_active": time.time()
        }
    session_data[session_hash]["last_active"] = time.time()
    return session_data[session_hash]

def save_feedback_to_hf(feedback_text, selected_persona, life_moment, request: gr.Request = None):
    sdata = get_session_data(request)
    latest = sdata["latest_interaction"]

    if not feedback_text or not feedback_text.strip():
        return "Please enter feedback text."
    if not latest["user"] or not latest["assistant"]:
        return "No interaction found. Please chat with the AI first."

    if not HF_TOKEN or not HF_DATASET_NAME:
        return "Error: HF_TOKEN or HF_DATASET_NAME not configured."

    data = {
        "timestamp": time.time(),
        "persona": selected_persona,
        "life_moment": life_moment if life_moment else "None",
        "user_message": latest["user"],
        "assistant_response": latest["assistant"],
        "expert_feedback": feedback_text
    }

    try:
        api = HfApi(token=HF_TOKEN)
        filename = "feedback.jsonl"

        # Ensure the dataset repo exists (create if it doesn't)
        try:
            api.create_repo(
                repo_id=HF_DATASET_NAME,
                repo_type="dataset",
                exist_ok=True,
                private=False
            )
        except Exception:
            pass  # Repo likely already exists

        # Download existing or start new
        temp_dir = tempfile.gettempdir()
        upload_path = os.path.join(temp_dir, f"feedback_{int(time.time())}.jsonl")

        try:
            # Force download latest version (bypass cache)
            local_path = hf_hub_download(
                repo_id=HF_DATASET_NAME,
                filename=filename,
                repo_type="dataset",
                token=HF_TOKEN,
                force_download=True
            )
            # Read existing content and write to temp file with new entry
            with open(local_path, "r", encoding="utf-8") as f:
                existing_content = f.read()
            with open(upload_path, "w", encoding="utf-8") as f:
                f.write(existing_content)
                if existing_content and not existing_content.endswith("\n"):
                    f.write("\n")
                f.write(json.dumps(data, ensure_ascii=False) + "\n")
        except Exception:
            # File doesn't exist yet on HuggingFace - create new
            with open(upload_path, "w", encoding="utf-8") as f:
                f.write(json.dumps(data, ensure_ascii=False) + "\n")

        api.upload_file(
            path_or_fileobj=upload_path,
            path_in_repo=filename,
            repo_id=HF_DATASET_NAME,
            repo_type="dataset"
        )
        return "Feedback submitted successfully!"
    except Exception as e:
        return f"Error uploading: {str(e)}"

def chat_wrapper(message: str, history: list, selected_model: str, life_moment: str, use_counselor: bool, request: gr.Request):
    sdata = get_session_data(request)
    sdata["current_dialogue"] = ""  # Clear state for new turn
    full_output = ""
    for output, dialogue in chat(message, history, selected_model, life_moment, use_counselor):
        if dialogue:
            sdata["current_dialogue"] = dialogue
        full_output = output
        yield output

    # Store latest once done
    sdata["latest_interaction"] = {"user": message, "assistant": full_output}

def chat(message: str, history: list, selected_model: str = "Hitchens 1", life_moment: str = "None", use_counselor: bool = False):
    print(f"\n{'='*60}")
    print(f"CHAT: New message at turn #{len(history)}, message='{message[:80]}...'")
    print(f"CHAT: history entries={len(history)}, model={selected_model}, life_moment={life_moment}")
    print(f"{'='*60}")

    state = {
        "message": message,
        "history": history,
        "route": "kantian",
        "response": "",
        "progress": "",
        "summary": "",
        "used_concepts": [],
        "selected_model": selected_model,
        "life_moment": life_moment,
        "use_counselor": use_counselor
    }

    config = {"recursion_limit": 100}
    output = ""
    debate_dialogue = ""

    try:
        for step in app.stream(state, config=config):
            for value in step.values():
                if isinstance(value, dict):
                    if "response" in value and value["response"]:
                        output = value["response"]
                    elif "progress" in value and value["progress"]:
                        output = value["progress"]
                    # Check if debate progress update
                    if "progress" in value and "**R1 - Proposition:**" in str(value.get("progress", "")):
                        debate_dialogue = value["progress"]
                    # Check if debate was triggered (final result)
                    elif "debate_r1" in value:
                        debate_dialogue = f"**R1 - Proposition:**\n{value['debate_r1']}\n\n**R2 - Internal Critique:**\n{value['debate_r2']}\n\n**R3 - Final Synthesis:**\n{value['debate_r3']}"
            yield output, debate_dialogue
    except Exception as e:
        print(f"FATAL ERROR in chat() at turn #{len(history)}: {type(e).__name__}: {e}")
        traceback.print_exc()
        output = f"Error: {type(e).__name__} - {str(e)[:200]}. Please reset the chat."
        yield output, debate_dialogue

    if not output.strip():
        output = "I am ready. Ask me a question."
        yield output, debate_dialogue

def get_dialogue(history, request: gr.Request):
    sdata = get_session_data(request)
    if not history:
        sdata["current_dialogue"] = ""
    return sdata["current_dialogue"]

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Philosophical Chat")
    gr.Markdown("Chat with a Hitchens persona")

    with gr.Row():
        model_dropdown = gr.Dropdown(
            choices=list(MODEL_OPTIONS.keys()),
            value="Hitchens 1",  # Default to Hitchens 1
            label="Select Persona",
            interactive=True
        )
        life_moment_dropdown = gr.Dropdown(
            choices=PERSONA_LIFE_MOMENTS["Hitchens 1"],
            value="None",
            label="Life Moment (Hitchens only)",
            interactive=True,
            visible=True  # Visible by default since Hitchens is selected
        )
        use_counselor_toggle = gr.Checkbox(
            label="Use Counselor Mode",
            value=USE_COUNSELOR_AS_DEFAULT,
            info="ON: Sociological Counselor prompts | OFF: Raw persona prompts",
            interactive=True
        )

    # Update life moment dropdown when persona changes
    def update_life_moments(persona):
        moments = PERSONA_LIFE_MOMENTS.get(persona, ["None"])
        is_hitchens_persona = is_hitchens(persona)
        return gr.update(choices=moments, value="None", visible=is_hitchens_persona)

    model_dropdown.change(
        fn=update_life_moments,
        inputs=[model_dropdown],
        outputs=[life_moment_dropdown]
    )

    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.ChatInterface(
                fn=chat_wrapper,
                additional_inputs=[model_dropdown, life_moment_dropdown, use_counselor_toggle],
                textbox=gr.Textbox(placeholder="Ask me anything...", interactive=True)
            )

        with gr.Column(scale=1):
            internal_dialogue = gr.Textbox(
                label="Internal Debate Dialogue",
                lines=20,
                interactive=False,
                placeholder="Debate rounds will appear here when triggered...",
                value=""
            )

    # Feedback Section
    with gr.Accordion("Expert Feedback (Fine-tuning Storage)", open=False):
        last_pair_display = gr.Markdown("No interaction yet. Chat with the AI to provide feedback.")
        feedback_input = gr.Textbox(label="Expert Critique", placeholder="What should the AI have said instead?", lines=3)
        with gr.Row():
            submit_feedback = gr.Button("Submit Feedback", variant="primary")
            skip_feedback = gr.Button("Skip/Clear")
        feedback_status = gr.Markdown("")

    def update_feedback_display(request: gr.Request = None):
        sdata = get_session_data(request)
        latest = sdata["latest_interaction"]
        if not latest["user"]:
            return "No interaction yet."
        return f"**User**: {latest['user']}\n\n**Assistant**: {latest['assistant']}"

    def clear_feedback():
        return "", ""

    def refresh_feedback_ui(request: gr.Request = None):
        return update_feedback_display(request), ""

    # Update feedback display after each turn
    chatbot.chatbot.change(
        refresh_feedback_ui,
        outputs=[last_pair_display, feedback_status]
    )

    submit_feedback.click(
        save_feedback_to_hf,
        inputs=[feedback_input, model_dropdown, life_moment_dropdown],
        outputs=feedback_status
    ).then(lambda: "", outputs=feedback_input)

    skip_feedback.click(clear_feedback, outputs=[feedback_input, feedback_status])

    # Update dialogue after each chat response
    chatbot.chatbot.change(
        get_dialogue,
        inputs=[chatbot.chatbot],
        outputs=internal_dialogue
    )

    reset_btn = gr.Button("Reset Chat", variant="secondary")

    def reset_all(request: gr.Request):
        sdata = get_session_data(request)
        sdata["current_dialogue"] = ""
        sdata["latest_interaction"] = {"user": "", "assistant": ""}
        return "", ""

    reset_btn.click(reset_all, outputs=[internal_dialogue, feedback_status]).then(None, js="() => window.location.reload()")

if __name__ == "__main__":
    print("Starting Philosophical Chat (Multi-Model)")
    print(f"Available models: {list(MODEL_OPTIONS.keys())}")
    for name, model in MODEL_OPTIONS.items():
        print(f"  {name}: {model}")
    demo.launch(share=False)