"""LangGraph agent — one branch per question class.

The graph is deliberately shallow. There is no general-purpose tool loop and no
"let the model decide what to do", because the class decides the tool path and
the class is assigned before generation starts. A loop that could route a
procedural question through the historical branch is the failure mode this
whole design exists to prevent.

    classify ─┬─ historical ─┐
              ├─ reference   ├─ generate → validate → (retry once) → answer
              ├─ procedural  │
              ├─ advisory   ─┘
              └─ unclear ────── clarify (no generation, no tools)

Each branch gathers its own evidence and hands a payload to the contract. The
prose model never sees a tool it was not given for its class.
"""

from __future__ import annotations

import json
import logging
from typing import Any, TypedDict

from langgraph.graph import END, StateGraph
from openai import AzureOpenAI

import classifier
import tools.equipment as equipment_tool
import tools.metrics as metrics
import tools.retrieval as retrieval
from config import settings
from contracts import ContractViolation, QuestionClass
from guardrails import enforce_contract

log = logging.getLogger("agent")


class State(TypedDict, total=False):
    question: str
    classification: classifier.Classification
    evidence: dict[str, Any]
    payload: dict[str, Any]
    answer: Any
    trace: Any


def _client() -> AzureOpenAI:
    cfg = settings()
    return AzureOpenAI(
        azure_endpoint=cfg.azure_openai_endpoint,
        api_key=cfg.azure_openai_api_key,
        api_version=cfg.azure_openai_api_version,
    )


def _embed(text: str, *, client: AzureOpenAI) -> list[float]:
    cfg = settings()
    return client.embeddings.create(model=cfg.embed_deployment, input=text).data[0].embedding


# ---------------------------------------------------------------------------
# System prompts.
#
# Byte-identical between calls so prompt caching applies - never interpolate
# the question, the evidence or the date into these. Everything variable goes
# in the user message.
#
# These prompts describe the job. They are NOT the safety mechanism; the
# contracts are. If a rule matters, it is in contracts.py as well as here, and
# if the two ever disagree, contracts.py wins by construction.
# ---------------------------------------------------------------------------

_BASE = """\
You write answers for operators at the Waterloo Road Pump Station, a three-pump
wastewater station. Be brief and concrete. Use the equipment names an operator
uses. Always state units, and state the time window any figure came from.
Never state a figure that is not in the evidence you were given.
Return JSON matching the schema you are given, and nothing else.
"""

HISTORICAL_PROMPT = _BASE + """\
You are answering a question about what happened. You have query results.
If the result set is empty, say plainly that no records were found for the
window - do not estimate, do not reason about what the number probably was.
"""

REFERENCE_PROMPT = _BASE + """\
You are answering a question about what something is or means. You have
document extracts and tag metadata. Cite every claim: document number,
revision, effective date. If nothing relevant was retrieved, say so.
"""

PROCEDURAL_PROMPT = _BASE + """\
You are IDENTIFYING a controlled procedure, not explaining it.

Give the procedure number, revision, effective date, title, the authorising
role, and where the controlled copy is. Quote prerequisites word for word into
prerequisites_verbatim. Write no steps, no paraphrase of steps, no summary of
what the procedure involves, and no advice about what to do first. An interlock
exists because someone assessed a hazard; a bypass procedure you reconstructed
is a safety document nobody approved.

If no controlled procedure was retrieved, say so and stop.
"""

ADVISORY_PROMPT = _BASE + """\
You are presenting EVIDENCE, not a recommendation.

Report what has actually been done: the rates used, how many operations that is
drawn from, the outcomes, when alarms occurred, and the documented limits with
their citations. Then defer explicitly to a competent person.

Do not state a recommended value. Do not say what is best, optimal or safest.
Do not offer a range as a disguised recommendation. "Best" depends on equipment
condition and concurrent operations you cannot see, and any number you present
as an answer will be typed into a control system by someone who trusts it.
"""

PROMPTS = {
    QuestionClass.HISTORICAL: HISTORICAL_PROMPT,
    QuestionClass.REFERENCE: REFERENCE_PROMPT,
    QuestionClass.PROCEDURAL: PROCEDURAL_PROMPT,
    QuestionClass.ADVISORY: ADVISORY_PROMPT,
}


# ---------------------------------------------------------------------------
# Evidence gathering — one node per class. No node can reach another's tools.
# ---------------------------------------------------------------------------


def _resolve_entities(state: State) -> tuple[str | None, list[str]]:
    """Resolve the first equipment entity the classifier found.

    Returns (equipment_id, ambiguities). An ambiguous term is surfaced, not
    guessed at - agent behaviour on ambiguity is to say what it matched.
    """
    named = (state["classification"].entities or {}).get("equipment") or []
    for term in named:
        matches = equipment_tool.resolve(str(term))
        equipment = [m for m in matches if m.kind == "equipment"]
        if len(equipment) == 1:
            return equipment[0].canonical_id, []
        if len(equipment) > 1:
            return None, [m.canonical_id for m in equipment]
    return None, []


def gather_historical(state: State) -> State:
    trace = state.get("trace")
    equipment_id, _ = _resolve_entities(state)
    result = metrics.run(metrics.alarm_detail(equipment_id=equipment_id, days=7), trace=trace)
    _, _, window_description = metrics.rolling_window(7)
    state["evidence"] = {
        "query": result.query,
        "rows": result.rows,
        "row_count": result.row_count,
        "time_window": {**result.time_window, "description": window_description},
        "used_fixture_data": result.used_fixture_data,
    }
    return state


def gather_reference(state: State) -> State:
    client = _client()
    equipment_id, _ = _resolve_entities(state)
    embedding = _embed(state["question"], client=client)
    chunks = retrieval.rerank(
        retrieval.search(embedding, top_k=8, equipment_id=equipment_id),
        state["question"],
    )
    tag_rows = equipment_tool.tags_for_equipment(equipment_id) if equipment_id else []
    state["evidence"] = {
        "chunks": retrieval.as_dicts(chunks),
        "citations": [c.citation() for c in chunks],
        "tags": tag_rows,
        "used_fixture_data": False,
    }
    return state


def gather_procedural(state: State) -> State:
    """Procedures only, live revisions only. Nothing else is in scope here."""
    client = _client()
    equipment_id, _ = _resolve_entities(state)
    embedding = _embed(state["question"], client=client)
    chunks = retrieval.find_procedure(
        embedding, state["question"], equipment_id=equipment_id
    )
    state["evidence"] = {
        "chunks": retrieval.as_dicts(chunks),
        "citations": [c.citation() for c in chunks],
        "used_fixture_data": False,
    }
    return state


def gather_advisory(state: State) -> State:
    """Both paths: what was done (Cube) and what is allowed (documents)."""
    trace = state.get("trace")
    client = _client()
    result = metrics.run(metrics.pump_down_evidence(days=30), trace=trace)
    _, _, window_description = metrics.rolling_window(30)
    embedding = _embed(state["question"], client=client)
    chunks = retrieval.rerank(
        retrieval.search(embedding, top_k=8, doc_type="design"), state["question"]
    )
    state["evidence"] = {
        "query": result.query,
        "rows": result.rows,
        "row_count": result.row_count,
        "time_window": {**result.time_window, "description": window_description},
        "chunks": retrieval.as_dicts(chunks),
        "citations": [c.citation() for c in chunks],
        "used_fixture_data": result.used_fixture_data,
    }
    return state


def gather_unclear(state: State) -> State:
    """No tools, no generation. The clarifying question is built from what the
    classifier said was missing, so it asks for something specific."""
    missing = state["classification"].missing_context or ["what you are asking about"]
    asked_for = ", ".join(missing)
    state["evidence"] = {}
    state["payload"] = {
        "question": state["question"],
        "question_class": QuestionClass.UNCLEAR,
        "answer": f"I need one more detail before I can answer: {asked_for}.",
        "clarifying_question": f"Could you tell me {asked_for}?",
        "candidate_interpretations": [],
        "used_fixture_data": False,
    }
    return state


# ---------------------------------------------------------------------------
# Generation and enforcement
# ---------------------------------------------------------------------------

_SCHEMA_HINTS: dict[QuestionClass, str] = {
    QuestionClass.HISTORICAL: (
        '{"answer": str}  - the figures, the window, and what the rows show'
    ),
    QuestionClass.REFERENCE: '{"answer": str}',
    QuestionClass.PROCEDURAL: (
        '{"answer": str, "procedure": {"doc_number": str, "title": str, '
        '"revision": str, "effective_date": "YYYY-MM-DD", '
        '"authorising_role": str, "controlled_copy_location": str}, '
        '"prerequisites_verbatim": [str]}'
    ),
    QuestionClass.ADVISORY: (
        '{"answer": str, "evidence": [{"description": str, "metric": str, '
        '"value": number, "unit": str, "sample_size": int}], '
        '"documented_limits": [], "deferral": str}'
    ),
}


def generate(state: State) -> State:
    """Generate prose, then enforce the contract. Retry once, then error."""
    klass = state["classification"].question_class
    if klass is QuestionClass.UNCLEAR:
        return state  # gather_unclear already built the payload

    cfg = settings()
    client = _client()
    trace = state.get("trace")
    evidence = state["evidence"]

    def call(attempt: int, previous: ContractViolation | None) -> dict[str, Any]:
        user = {
            "question": state["question"],
            "evidence": evidence,
            "return_schema": _SCHEMA_HINTS[klass],
        }
        if previous is not None:
            # Tell it what it broke. One retry only - a model that fails a
            # safety contract twice is not going to be argued into compliance.
            user["previous_attempt_rejected"] = {
                "rule": previous.rule,
                "detail": previous.detail,
            }
        response = client.chat.completions.create(
            model=cfg.chat_deployment,
            messages=[
                {"role": "system", "content": PROMPTS[klass]},
                {"role": "user", "content": json.dumps(user, default=str)},
            ],
            temperature=0.1,
            max_tokens=cfg.max_output_tokens,
            response_format={"type": "json_object"},
        )
        generated = json.loads(response.choices[0].message.content)

        # The model supplies prose and its own structured fields. Everything
        # factual - rows, counts, citations, the fixture flag - is attached
        # here from the evidence, so the model cannot alter it.
        payload: dict[str, Any] = {
            "question": state["question"],
            "question_class": klass,
            "answer": generated.get("answer", ""),
            "used_fixture_data": evidence.get("used_fixture_data", False),
        }
        if klass is QuestionClass.HISTORICAL:
            payload.update(
                query=evidence["query"],
                rows=evidence["rows"],
                row_count=evidence["row_count"],
                time_window=evidence["time_window"],
                citations=evidence.get("citations", []),
            )
        elif klass is QuestionClass.REFERENCE:
            payload.update(
                citations=evidence.get("citations", []),
                tags_referenced=[t["tag_id"] for t in evidence.get("tags", [])],
            )
        elif klass is QuestionClass.PROCEDURAL:
            payload.update(
                procedure=generated.get("procedure"),
                prerequisites_verbatim=generated.get("prerequisites_verbatim", []),
                steps_provided=False,
                citations=evidence.get("citations", []),
            )
        elif klass is QuestionClass.ADVISORY:
            payload.update(
                evidence=generated.get("evidence", []),
                documented_limits=generated.get("documented_limits", []),
                recommendation_given=False,
                deferral=generated.get("deferral", ""),
                citations=evidence.get("citations", []),
            )
        return payload

    result = enforce_contract(call, klass, trace=trace)
    state["answer"] = result.answer
    state["payload"] = result.answer.model_dump()
    return state


def validate_unclear(state: State) -> State:
    from contracts import validate_answer

    if state["classification"].question_class is QuestionClass.UNCLEAR:
        state["answer"] = validate_answer(state["payload"], QuestionClass.UNCLEAR)
    return state


# ---------------------------------------------------------------------------
# Graph
# ---------------------------------------------------------------------------


def classify_node(state: State) -> State:
    state["classification"] = classifier.classify(
        state["question"], client=_client(), trace=state.get("trace")
    )
    return state


def route(state: State) -> str:
    return state["classification"].question_class.value


def build_graph():
    graph = StateGraph(State)
    graph.add_node("classify", classify_node)
    graph.add_node("historical", gather_historical)
    graph.add_node("reference", gather_reference)
    graph.add_node("procedural", gather_procedural)
    graph.add_node("advisory", gather_advisory)
    graph.add_node("unclear", gather_unclear)
    graph.add_node("generate", generate)
    graph.add_node("finalise", validate_unclear)

    graph.set_entry_point("classify")
    graph.add_conditional_edges(
        "classify",
        route,
        {
            "historical": "historical",
            "reference": "reference",
            "procedural": "procedural",
            "advisory": "advisory",
            "unclear": "unclear",
        },
    )
    for node in ("historical", "reference", "procedural", "advisory", "unclear"):
        graph.add_edge(node, "generate")
    graph.add_edge("generate", "finalise")
    graph.add_edge("finalise", END)
    return graph.compile()


_GRAPH = None


def answer(question: str, *, trace=None):
    """Answer one question. Raises ContractViolation if the contract cannot be
    met - the caller returns an error, never a partial answer."""
    global _GRAPH
    if _GRAPH is None:
        _GRAPH = build_graph()
    final = _GRAPH.invoke({"question": question, "trace": trace})
    return final["answer"], final["classification"]