"""LangGraph node implementations for the VeriScite audit pipeline."""

import json
import os
from langchain_groq import ChatGroq
from langgraph.config import get_stream_writer

from app.graph.state import GraphState, ClaimCitation
from app.tools import claim_client, llm_verifier, semantic_scholar


def _writer():
    """Safe wrapper around LangGraph's custom stream writer.

    get_stream_writer() raises RuntimeError when called outside an actual
    graph run (e.g. our isolated node tests, which call node functions
    directly rather than through graph.ainvoke/astream). Falls back to a
    no-op so those existing tests keep working unchanged; real graph runs
    get live progress events via stream_mode="custom".
    """
    try:
        return get_stream_writer()
    except RuntimeError:
        return lambda _event: None

EXTRACT_PROMPT = """This text contains a scientific paper's body (with inline citation \
markers like [12] or (Smith et al., 2023)) followed by its bibliography/reference list.

For each claim in the body that is attributed to a citation, extract:
  - claim: the factual statement being made
  - citation_marker: the inline marker exactly as it appears in the body, e.g. "[12]"
  - reference_string: the FULL bibliography entry that marker resolves to (authors,
    year, title) — look this up in the reference list, do not guess or invent one.
    If the marker cannot be resolved to a reference list entry, skip that claim.

Respond as a JSON list only, no other text:
[{{"claim": "...", "citation_marker": "...", "reference_string": "..."}}, ...]

Text:
{text}"""

AGREEMENT_THRESHOLD = 0.5  # min confidence gap tolerated before escalation

MAX_AGENT_ITERATIONS = 3  # ReAct-loop safety cap — avoids runaway free-tier usage

AGENT_SYSTEM_PROMPT = """You are resolving a disagreement between two independent \
methods that assessed whether a piece of evidence supports a scientific claim.

Claim: {claim}
Evidence currently available: {evidence}

Assessment A (fine-tuned NLI model): {label_a} (confidence {confidence_a})
Assessment B (zero-shot LLM verifier): {label_b} (confidence {confidence_b})
Assessment B's reasoning: {reasoning_b}

You have tools available to investigate further before concluding. Use them if \
they would genuinely help; call `conclude` as soon as you have a well-supported \
answer. Do not call more tools than necessary. If you cannot resolve the \
disagreement with more evidence, concluding NOT_ENOUGH_INFO is a legitimate, \
complete answer — it is not a failure, and is preferable to guessing."""

AGENT_TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "retry_with_query",
            "description": ("Search Semantic Scholar again with a different query, "
                             "in case the currently available evidence missed the "
                             "relevant passage. Write your own search query."),
            "parameters": {
                "type": "object",
                "properties": {"query": {"type": "string", "description": "New search query"}},
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "request_second_opinion",
            "description": ("Get an independent, blind re-verification of the claim "
                             "against the current evidence (no prior verdicts shown "
                             "to it, to avoid anchoring bias)."),
            "parameters": {"type": "object", "properties": {}},
        },
    },
    {
        "type": "function",
        "function": {
            "name": "conclude",
            "description": "Give your final, complete answer. This ends the investigation.",
            "parameters": {
                "type": "object",
                "properties": {
                    "label": {"type": "string", "enum": ["SUPPORT", "NOT_ENOUGH_INFO", "CONTRADICT"]},
                    "confidence": {"type": "number"},
                    "reasoning": {"type": "string"},
                },
                "required": ["label", "confidence", "reasoning"],
            },
        },
    },
]


def _llm():
    return ChatGroq(model="openai/gpt-oss-20b", temperature=0.0,
                     api_key=os.environ.get("GROQ_API_KEY"))


def _planner_llm():
    # Deliberately a different, actively-supported model from the verifier
    # (gpt-oss-20b) — a same-model critic tends to rubber-stamp its own kind
    # of reasoning (Panickssery et al., 2024). llama-3.3-70b-versatile was
    # deprecated by Groq (announced 2026-06-17); gpt-oss-120b is Groq's own
    # recommended replacement and remains actively supported.
    return ChatGroq(model="openai/gpt-oss-120b", temperature=0.0,
                     api_key=os.environ.get("GROQ_API_KEY")).bind_tools(AGENT_TOOLS)


async def extract_claims(state: GraphState) -> GraphState:
    writer = _writer()
    writer({"node": "extract_claims", "status": "start"})
    prompt = EXTRACT_PROMPT.format(text=state["source_text"])
    response = await _llm().ainvoke(prompt)
    try:
        pairs = json.loads(response.content)
    except json.JSONDecodeError:
        pairs = []
    claims: list[ClaimCitation] = [
        {"claim": p["claim"], "citation_marker": p["citation_marker"],
         "reference_string": p.get("reference_string", ""),
         "resolved_paper_id": None, "evidence_text": None}
        for p in pairs
        if p.get("reference_string")  # skip claims whose citation couldn't be resolved
    ]
    writer({"node": "extract_claims", "status": "done", "n_claims": len(claims)})
    return {**state, "claims": claims, "audits": [], "current_index": 0}


async def fetch_citation(state: GraphState) -> GraphState:
    writer = _writer()
    claims = state["claims"]
    idx = state["current_index"]
    cc = claims[idx]
    writer({"node": "fetch_citation", "status": "start", "claim": cc["claim"]})
    s2_key = os.environ.get("S2_API_KEY")
    query = semantic_scholar.extract_query(cc["reference_string"])
    paper = await semantic_scholar.search_paper(query, api_key=s2_key)
    if paper and paper.get("abstract"):
        cc["resolved_paper_id"] = paper.get("paperId")
        cc["evidence_text"] = paper.get("abstract")
        writer({"node": "fetch_citation", "status": "done", "resolved": True,
                "title": paper.get("title")})
    else:
        # Either no match, or matched but no abstract indexed (common for older/
        # non-open-access papers) — both are "citation not usable", not a partial
        # success. Leave resolved_paper_id=None so downstream nodes can detect this.
        cc["resolved_paper_id"] = None
        cc["evidence_text"] = None
        writer({"node": "fetch_citation", "status": "done", "resolved": False})
    claims[idx] = cc
    return {**state, "claims": claims}


async def verify_dual(state: GraphState) -> GraphState:
    idx = state["current_index"]
    cc = state["claims"][idx]
    claim, evidence = cc["claim"], cc["evidence_text"] or ""
    writer = _writer()
    writer({"node": "verify_dual", "status": "start"})

    deberta_res = await claim_client.analyze(claim, evidence)
    winner = deberta_res["winner"]  # {"sentence", "label", "confidence", "sentence_index"}
    llm_res = await llm_verifier.verify(claim, evidence)

    agree = winner["label"] == llm_res.get("label")
    deberta_verdict = {"label": winner["label"], "confidence": winner["confidence"], "source": "deberta"}
    writer({"node": "verify_dual", "status": "done", "agree": agree,
            "deberta_label": winner["label"], "llm_label": llm_res.get("label")})
    audit = {
        "claim_citation": cc,
        "winner_sentence": winner["sentence"],
        "attribution_available": deberta_res.get("attribution_available", False),
        "deberta_verdict": deberta_verdict,
        "llm_verdict": {"label": llm_res.get("label"),
                         "confidence": llm_res.get("confidence", 0.0),
                         "source": "llm"},
        "agreement": agree,
        "escalated": False,
        "escalation_trace": None,
        "final_verdict": deberta_verdict if agree else None,
        "attribution": None,
        "resolution_note": None,
    }
    return {**state, "audits": state["audits"] + [audit]}


def route_after_fetch(state: GraphState) -> str:
    """Conditional edge: skip verification entirely if the citation could not
    be resolved to usable evidence text (search miss, or resolved paper had
    no abstract indexed) — sending empty evidence to /analyze causes a 422
    from clAIm's backend rather than a meaningful verdict.
    """
    idx = state["current_index"]
    cc = state["claims"][idx]
    return "verify" if cc.get("evidence_text") else "unresolved"


async def handle_unresolved_citation(state: GraphState) -> GraphState:
    idx = state["current_index"]
    cc = state["claims"][idx]
    audit = {
        "claim_citation": cc,
        "winner_sentence": None,
        "attribution_available": False,
        "deberta_verdict": None,
        "llm_verdict": None,
        "agreement": None,
        "escalated": False,
        "escalation_trace": None,
        "final_verdict": None,
        "attribution": None,
        "resolution_note": ("Citation could not be resolved to usable evidence — "
                             "either no search match or resolved paper had no "
                             "abstract indexed on Semantic Scholar."),
    }
    return {**state, "audits": state["audits"] + [audit]}


def route_after_verify(state: GraphState) -> str:
    """Conditional edge: escalate on disagreement, else proceed to explain."""
    last_audit = state["audits"][-1]
    return "escalate" if not last_audit["agreement"] else "explain"


async def escalate(state: GraphState) -> GraphState:
    """On disagreement: a planner model (gpt-oss-120b, distinct from the
    gpt-oss-20b verifier) autonomously decides how to resolve it -- it can
    retry retrieval with its own reformulated query, request a blind second
    opinion, or conclude directly, in whatever order and however many times
    (up to MAX_AGENT_ITERATIONS) it judges necessary. This is a ReAct-style
    reason-act-observe loop: the control flow is decided by the model at
    each step, not by fixed code. See docs/dev_log.md for design rationale.
    """
    idx = state["current_index"]
    cc = state["claims"][idx]
    audits = state["audits"]
    last = audits[-1]
    last["escalated"] = True
    writer = _writer()
    writer({"node": "escalate", "status": "start", "reason": "verifiers disagreed"})

    from langchain_core.messages import SystemMessage, ToolMessage

    current_evidence = cc["evidence_text"] or ""
    system_prompt = AGENT_SYSTEM_PROMPT.format(
        claim=cc["claim"],
        evidence=current_evidence,
        label_a=last["deberta_verdict"]["label"],
        confidence_a=last["deberta_verdict"]["confidence"],
        label_b=last["llm_verdict"]["label"],
        confidence_b=last["llm_verdict"]["confidence"],
        reasoning_b="",  # llm_verifier.verify doesn't currently surface reasoning to verify_dual
    )
    messages = [SystemMessage(content=system_prompt)]
    trace = []
    planner = _planner_llm()
    s2_key = os.environ.get("S2_API_KEY")

    for iteration in range(MAX_AGENT_ITERATIONS):
        try:
            response = await planner.ainvoke(messages)
        except Exception as exc:
            trace.append({"action": "planner_error", "input": None, "observation": str(exc)})
            writer({"node": "escalate", "status": "action", "action": "planner_error",
                    "iteration": iteration})
            break

        messages.append(response)

        if not response.tool_calls:
            messages.append(SystemMessage(content="Please call the `conclude` tool with your final answer."))
            continue

        tool_call = response.tool_calls[0]  # one action at a time, per ReAct
        name, args = tool_call["name"], tool_call["args"]
        writer({"node": "escalate", "status": "action", "action": name,
                "input": args, "iteration": iteration})

        if name == "conclude":
            last["final_verdict"] = {"label": args["label"], "confidence": args.get("confidence", 0.0),
                                      "source": "agent"}
            trace.append({"action": "conclude", "input": args, "observation": None})
            last["escalation_trace"] = trace
            writer({"node": "escalate", "status": "done", "final_verdict": last["final_verdict"]})
            return {**state, "audits": audits}

        elif name == "retry_with_query":
            paper = await semantic_scholar.search_paper(args["query"], api_key=s2_key)
            if paper and paper.get("abstract"):
                current_evidence = paper["abstract"]
                deberta_res = await claim_client.analyze(cc["claim"], current_evidence)
                winner = deberta_res["winner"]
                new_llm_res = await llm_verifier.verify(cc["claim"], current_evidence)
                observation = (f"New evidence found: \"{current_evidence[:300]}\". "
                                f"Re-verified: NLI model says {winner['label']} "
                                f"({winner['confidence']:.2f}), LLM verifier says "
                                f"{new_llm_res.get('label')} ({new_llm_res.get('confidence', 0):.2f}).")
                last["winner_sentence"] = winner["sentence"]
                last["attribution_available"] = deberta_res.get("attribution_available", False)
                last["deberta_verdict"] = {"label": winner["label"], "confidence": winner["confidence"], "source": "deberta"}
                last["llm_verdict"] = {"label": new_llm_res.get("label"), "confidence": new_llm_res.get("confidence", 0.0), "source": "llm"}
            else:
                observation = "No usable evidence found for that query (no match or no abstract indexed)."
            trace.append({"action": "retry_with_query", "input": args, "observation": observation})
            writer({"node": "escalate", "status": "observation", "action": "retry_with_query",
                    "observation": observation})
            messages.append(ToolMessage(content=observation, tool_call_id=tool_call["id"]))

        elif name == "request_second_opinion":
            second = await llm_verifier.verify(cc["claim"], current_evidence)  # blind -- no prior verdicts in this call
            observation = (f"Second opinion (blind, independent): {second.get('label')} "
                            f"(confidence {second.get('confidence', 0):.2f}). "
                            f"Reasoning: {second.get('reasoning', '')}")
            trace.append({"action": "request_second_opinion", "input": {}, "observation": observation})
            writer({"node": "escalate", "status": "observation", "action": "request_second_opinion",
                    "observation": observation})
            messages.append(ToolMessage(content=observation, tool_call_id=tool_call["id"]))

    # Iteration cap hit (or planner errored above) without a `conclude` call --
    # force a final answer rather than leaving the claim unresolved. Per
    # ReAct/Reflexion literature, a calibrated NOT_ENOUGH_INFO is a legitimate
    # complete answer, not a failure.
    writer({"node": "escalate", "status": "action", "action": "forced_conclude", "iteration": MAX_AGENT_ITERATIONS})
    messages.append(SystemMessage(content="You must conclude now with your best answer, "
                                           "using only the `conclude` tool."))
    try:
        final_response = await planner.ainvoke(messages)
    except Exception as exc:
        trace.append({"action": "planner_error_on_forced_conclude", "input": None, "observation": str(exc)})
        final_response = None

    if final_response and final_response.tool_calls and final_response.tool_calls[0]["name"] == "conclude":
        args = final_response.tool_calls[0]["args"]
        last["final_verdict"] = {"label": args["label"], "confidence": args.get("confidence", 0.0), "source": "agent"}
        trace.append({"action": "conclude (forced at iteration cap)", "input": args, "observation": None})
    else:
        # Model never produced a valid conclude call, either it declined to
        # or Groq's parser failed on it. Fall back to NOT_ENOUGH_INFO rather
        # than guess, consistent with the "refusal over fabrication" principle,
        # and rather than crash the request.
        last["final_verdict"] = {"label": "NOT_ENOUGH_INFO", "confidence": 0.0, "source": "agent_fallback"}
        trace.append({"action": "forced_fallback", "input": None, "observation": "Model did not call conclude within iteration cap."})

    last["escalation_trace"] = trace
    writer({"node": "escalate", "status": "done", "final_verdict": last["final_verdict"]})
    return {**state, "audits": audits}


async def explain(state: GraphState) -> GraphState:
    idx = state["current_index"]
    cc = state["claims"][idx]
    audits = state["audits"]
    last = audits[-1]
    writer = _writer()
    writer({"node": "explain", "status": "start"})
    if last["final_verdict"] is None:
        # Should not normally happen -- verify_dual sets it on agreement,
        # escalate always sets it via conclude/forced-conclude/fallback.
        last["final_verdict"] = last["deberta_verdict"]

    if last["attribution_available"]:
        label_id = claim_client.LABEL2ID[last["final_verdict"]["label"]]
        last["attribution"] = await claim_client.attribute(
            cc["claim"], last["winner_sentence"], label_id
        )
    else:
        last["attribution"] = None  # winner was NOT_ENOUGH_INFO — no attribution to show

    writer({"node": "explain", "status": "done", "final_verdict": last["final_verdict"]})
    return {**state, "audits": audits}


def advance_or_report(state: GraphState) -> str:
    """Conditional edge: loop to next claim, or finish and build report."""
    next_index = state["current_index"] + 1
    return "next_claim" if next_index < len(state["claims"]) else "report"


async def next_claim(state: GraphState) -> GraphState:
    return {**state, "current_index": state["current_index"] + 1}


async def build_report(state: GraphState) -> GraphState:
    writer = _writer()
    audits = state["audits"]
    verified = [a for a in audits if a["agreement"] is not None]  # excludes unresolved citations
    agreement_rate = (sum(a["agreement"] for a in verified) / len(verified)) if verified else 0.0
    n_escalated = sum(a["escalated"] for a in audits)
    escalated_audits = [a for a in audits if a["escalated"]]
    avg_agent_iterations = (
        sum(len(a.get("escalation_trace") or []) for a in escalated_audits) / len(escalated_audits)
        if escalated_audits else 0.0
    )
    n_unresolved = sum(1 for a in audits if a.get("resolution_note"))
    report = {
        "n_claims": len(audits),
        "n_unresolved_citations": n_unresolved,
        "initial_agreement_rate": agreement_rate,  # computed over resolved claims only
        "n_escalated": n_escalated,
        "avg_agent_iterations_when_escalated": avg_agent_iterations,
        "claims": audits,
    }
    writer({"node": "build_report", "status": "done"})
    return {**state, "report": report}