# EpiRAG - agents.py
# Multi-agent swarm debate engine with real-time SSE callbacks (v2).
#
# Architecture:
#   Round 1    - all debate agents respond in parallel (fast, independent views)
#   Rounds 2+  - agents respond sequentially, each sees full transcript before replying
#   Final      - Epsilon synthesizes, then Zeta audits citations only
#
# Agent roster (3 debate + 2 support):
#   Alpha   (llama-3.3-70b-versatile)             - Skeptic
#   Beta    (openai/gpt-oss-120b)                 - Deep Reasoner
#   Gamma   (qwen/qwen3.6-27b)                   - Pattern Connector
#   Epsilon (openai/gpt-oss-20b)                  - Synthesizer
#   Zeta    (llama-3.1-8b-instant)                - Citation Auditor

import concurrent.futures
from groq import Groq

AGENTS = [
    {
        "name":        "Alpha",
        "model":       "llama-3.3-70b-versatile",
        "provider":    "groq",
        "client_type": "groq",
        "color":       "red",
        "personality": (
            "You are Agent Alpha - a ruthless Skeptic. "
            "Challenge every claim aggressively. Demand evidence from the sources. "
            "Point out what is NOT supported. Be blunt and relentless. "
            "After every factual statement you accept, cite the source: [Source: <name>]."
        )
    },
    {
        "name":        "Beta",
        "model":       "openai/gpt-oss-120b",
        "provider":    "groq",
        "client_type": "groq",
        "color":       "yellow",
        "personality": (
            "You are Agent Beta - a Deep Reasoner. "
            "Think step by step. Show your chain of thought. "
            "Identify hidden assumptions and logical gaps the other agents miss. "
            "Precision and thoroughness over speed. "
            "After every factual statement, cite the source: [Source: <name>]."
        )
    },
    {
        "name":        "Gamma",
        "model":       "qwen/qwen3.6-27b",
        "provider":    "groq",
        "client_type": "groq",
        "color":       "green",
        "personality": (
            "You are Agent Gamma - a Pattern Connector. "
            "Find non-obvious connections between different sources. "
            "Look for relationships and synthesis opportunities others miss. "
            "Think laterally and creatively. "
            "After every factual statement, cite the source: [Source: <name>]."
        )
    },

    {
        "name":        "Epsilon",
        "model":       "openai/gpt-oss-20b",
        "provider":    "groq",
        "client_type": "groq",
        "color":       "blue",
        "personality": (
            "You are Agent Epsilon - the Synthesizer. "
            "Read all agents' arguments and produce a final authoritative answer. "
            "CRITICAL RULE: If the corpus excerpts don't contain the answer but the topic is "
            "within epidemic modeling, network science, or mathematical epidemiology - "
            "synthesize the answer from your own training knowledge. Label such facts as "
            "[General Knowledge] instead of a paper name. Never say 'cannot answer' for "
            "in-domain questions. For corpus-sourced facts use [Source: name]. "
            "Use LaTeX for all mathematical expressions (e.g. $R_0$, $\\beta$)."
        )
    },
    {
        "name":        "Zeta",
        "model":       "llama-3.1-8b-instant",
        "provider":    "groq",
        "client_type": "groq",
        "color":       "orange",
        "personality": (
            "You are Agent Zeta - the Citation Auditor. "
            "You receive a synthesized answer and the original source excerpts. "
            "OUTPUT FORMAT (strict): reproduce the answer text exactly as-is, making ONLY "
            "small inline tweaks where a [Source: name] tag cites a source that is NOT "
            "in the provided excerpts. In that case, change the tag to [Source: not verified]. "
            "NEVER output an 'Audit results:' section, bullet list of audit findings, or "
            "any commentary outside the answer text itself. "
            "Claims with [General Knowledge] tags - leave completely untouched. "
            "Claims with NO source tag - leave completely untouched. "
            "Just output the corrected answer and nothing else."
        )
    },
]

MAX_ROUNDS         = 2
MAX_TOKENS_AGENT   = 280
MAX_TOKENS_SYNTH   = 900
MAX_TOKENS_ZETA    = 800   # raised - Zeta was truncating mid-sentence at 500
TIMEOUT_SECONDS    = 25
CONTEXT_LIMIT      = 2500   # chars fed to synthesizer/zeta
AGENT_CONTEXT_LIMIT = 1200  # chars fed to individual debate agents (keep prompts short)

DOMAIN_GUARD = """
SCOPE: EpiRAG - strictly epidemic modeling, network science, mathematical epidemiology.
If the question is completely off-topic (not related to these fields), say so and stop.
If the question IS in-domain but the provided excerpts lack the answer:
  - Use your training knowledge to answer.
  - Label each such fact as [General Knowledge] instead of inventing source names.
  - Do NOT fabricate paper titles or authors.
  - Never say 'I cannot answer' for in-domain questions - always try.
"""


def _call_agent(agent, messages, groq_key, hf_token=None, max_tokens=MAX_TOKENS_AGENT):
    """Call a Groq-hosted agent. hf_token kept for signature compatibility."""
    try:
        client = Groq(api_key=groq_key)
        resp = client.chat.completions.create(
            model=agent["model"], messages=messages,
            temperature=0.7, max_tokens=max_tokens
        )
        return resp.choices[0].message.content.strip()
    except Exception as e:
        return f"[{agent['name']} error: {str(e)[:100]}]"


def _round1_msgs(agent, question, context):
    ctx = context[:AGENT_CONTEXT_LIMIT] + "..." if len(context) > AGENT_CONTEXT_LIMIT else context
    return [
        {"role": "system", "content": f"{DOMAIN_GUARD}\n\n{agent['personality']}"},
        {"role": "user",   "content": (
            f"Context from research papers/web:\n\n{ctx}\n\n---\n\n"
            f"Question: {question}\n\n"
            f"Answer concisely based on context. Cite sources after claims using [Source: name]. Stay in character."
        )}
    ]


def _round_n_msgs(agent, question, context, full_transcript: str):
    """
    Rounds 2+: each agent receives the transcript of all previous rounds
    (sequential mode) for richer context before responding.
    """
    ctx = context[:AGENT_CONTEXT_LIMIT] + "..." if len(context) > AGENT_CONTEXT_LIMIT else context
    # Cap transcript too - agents only need the gist, not the full wall of text
    transcript_cap = full_transcript[-2000:] if len(full_transcript) > 2000 else full_transcript
    return [
        {"role": "system", "content": f"{DOMAIN_GUARD}\n\n{agent['personality']}"},
        {"role": "user",   "content": (
            f"Context:\n\n{ctx}\n\nQuestion: {question}\n\n"
            f"=== DEBATE TRANSCRIPT SO FAR ===\n\n{transcript_cap}\n\n"
            f"=== YOUR TURN ===\n\n"
            f"Refine your position. Where do you agree or disagree? Be concise and cite sources."
        )}
    ]


def _synth_msgs(question, context, all_rounds):
    transcript = ""
    for i, rnd in enumerate(all_rounds, 1):
        transcript += f"\n\n{'='*40}\nROUND {i}\n{'='*40}\n"
        for name, ans in rnd.items():
            transcript += f"\n-- {name} --\n{ans}\n"
    ctx = context[:CONTEXT_LIMIT] + "..." if len(context) > CONTEXT_LIMIT else context
    return [
        {"role": "system", "content": (
            f"{DOMAIN_GUARD}\n\nYou are the Synthesizer. "
            "Produce the single best final answer following these rules STRICTLY:\n"
            "1. For facts taken directly from the provided context excerpts, cite with [Source: name].\n"
            "2. For established domain knowledge (textbook definitions, well-known models, "
            "   standard equations) that is NOT in the excerpts - write it WITHOUT any [Source:] tag. "
            "   Do NOT invent source names. Do NOT write [Source: ...] for things you know from training.\n"
            "3. Resolve agent disagreements using the strongest evidence from excerpts.\n"
            "4. Use LaTeX for all math (e.g. $R_0$, $\\beta$, $$\\frac{dI}{dt} = \\beta SI$$).\n"
            "5. Be concise and structured. End with: CONFIDENCE: HIGH / MEDIUM / LOW"
        )},
        {"role": "user", "content": (
            f"Context excerpts from corpus:\n\n{ctx}\n\n---\n\n"
            f"Question: {question}\n\n---\n\n"
            f"Debate transcript:{transcript}\n\n---\n\n"
            f"Produce the final synthesized answer. Only use [Source: name] for things "
            f"that appear in the context excerpts above."
        )}
    ]


def _factcheck_msgs(question, context, synthesis):
    ctx = context[:CONTEXT_LIMIT] + "..." if len(context) > CONTEXT_LIMIT else context
    # List the excerpt source names so Zeta knows what's available
    import re as _re
    source_names = list(dict.fromkeys(_re.findall(r'\[(?:LOCAL|WEB)\]\s*-\s*([^\[\]\(]+?)(?:\s*\(relevance)', ctx)))
    sources_list = ", ".join(source_names) if source_names else "(see excerpts above)"
    return [
        {"role": "system", "content": f"{DOMAIN_GUARD}\n\n{_get_zeta()['personality']}"},
        {"role": "user",   "content": (
            f"Available source excerpts (source names: {sources_list}):\n\n{ctx}\n\n---\n\n"
            f"Question: {question}\n\n---\n\n"
            f"Synthesized answer to audit:\n\n{synthesis}\n\n---\n\n"
            f"Check only explicit [Source: name] citations. "
            f"Uncited claims (no [Source:] tag) are established domain knowledge - leave them alone. "
            f"Output the answer with only small inline corrections where a cited source name "
            f"is not in the list above."
        )}
    ]


def _get_zeta():
    return next(a for a in AGENTS if a["name"] == "Zeta")


def _converged(answers):
    agree = ["i agree", "correct", "you're right", "i concur",
             "well said", "exactly", "this is accurate", "that's right"]
    hits  = sum(1 for a in answers.values()
                if any(p in a.lower() for p in agree))
    return hits >= len(answers) * 0.5


def _build_transcript(all_rounds: list[dict]) -> str:
    """Build a full text transcript of all rounds completed so far."""
    lines = []
    for i, rnd in enumerate(all_rounds, 1):
        lines.append(f"=== Round {i} ===")
        for name, ans in rnd.items():
            lines.append(f"-- {name} --\n{ans}")
        lines.append("")
    return "\n".join(lines)


def run_debate(question, context, groq_key, hf_token, callback=None):
    """
    Run the full multi-agent swarm debate (v2).

    Flow:
      Round 1    - parallel (independent first thoughts)
      Rounds 2-N - sequential (each agent sees full transcript before responding)
      Synthesis  - Epsilon synthesizes
      Fact-check - Zeta verifies citations

    callback(event: dict) is called after each agent responds, enabling SSE streaming.
    event shapes:
      {"type": "agent_done",   "round": int, "name": str, "color": str, "text": str}
      {"type": "round_start",  "round": int}
      {"type": "synthesizing"}
      {"type": "factchecking"}
      {"type": "done",         "consensus": bool, "rounds": int}
    Returns:
      {"final_answer", "debate_rounds", "consensus", "rounds_run", "agent_count"}
    """
    def emit(event):
        if callback:
            callback(event)

    debate_agents  = [a for a in AGENTS if a["name"] not in ("Epsilon", "Zeta")]
    synthesizer    = next(a for a in AGENTS if a["name"] == "Epsilon")
    fact_checker   = _get_zeta()
    agent_colors   = {a["name"]: a["color"] for a in AGENTS}
    debate_rounds  = []

    # -- Round 1: Parallel
    emit({"type": "round_start", "round": 1})
    round1 = {}

    with concurrent.futures.ThreadPoolExecutor(max_workers=len(debate_agents)) as ex:
        futures = {
            ex.submit(_call_agent, agent,
                      _round1_msgs(agent, question, context),
                      groq_key, hf_token): agent
            for agent in debate_agents
        }
        for future in concurrent.futures.as_completed(futures, timeout=TIMEOUT_SECONDS * 2):
            agent = futures[future]
            try:
                answer = future.result(timeout=TIMEOUT_SECONDS)
            except Exception as e:
                answer = f"[{agent['name']} timed out: {e}]"
            round1[agent["name"]] = answer
            emit({"type": "agent_done", "round": 1,
                  "name": agent["name"], "color": agent_colors[agent["name"]],
                  "text": answer})

    debate_rounds.append(round1)
    consensus    = _converged(round1)
    rounds_run   = 1

    # -- Rounds 2+: Sequential (each agent sees full transcript)
    while not consensus and rounds_run < MAX_ROUNDS:
        rounds_run += 1
        emit({"type": "round_start", "round": rounds_run})

        full_transcript = _build_transcript(debate_rounds)
        next_round      = {}

        for agent in debate_agents:
            msgs   = _round_n_msgs(agent, question, context, full_transcript)
            answer = _call_agent(agent, msgs, groq_key, hf_token)
            next_round[agent["name"]] = answer
            emit({"type": "agent_done", "round": rounds_run,
                  "name": agent["name"], "color": agent_colors[agent["name"]],
                  "text": answer})
            # Update transcript between agents so each sees the others' answers
            debate_rounds_temp = debate_rounds + [next_round]
            full_transcript    = _build_transcript(debate_rounds_temp)

        debate_rounds.append(next_round)
        consensus = _converged(next_round)

    # -- Synthesis
    emit({"type": "synthesizing"})
    ctx_trunc = context[:4000] if len(context) > 4000 else context
    synthesis = _call_agent(
        synthesizer,
        _synth_msgs(question, ctx_trunc, debate_rounds),
        groq_key, hf_token,
        max_tokens=MAX_TOKENS_SYNTH
    )

    # -- Fact-check
    emit({"type": "factchecking"})
    final = _call_agent(
        fact_checker,
        _factcheck_msgs(question, ctx_trunc, synthesis),
        groq_key, hf_token,
        max_tokens=MAX_TOKENS_ZETA
    )

    emit({"type": "done", "consensus": consensus, "rounds": rounds_run})

    return {
        "final_answer":  final,
        "debate_rounds": debate_rounds,
        "consensus":     consensus,
        "rounds_run":    rounds_run,
        "agent_count":   len(debate_agents)
    }