import json as _json
import logging
import os
import re as _re
import time
import threading
import requests as _requests
from collections import defaultdict
from pathlib import Path

from elasticsearch import Elasticsearch
from langchain_core.documents import Document
from langchain_core.messages import AIMessage, SystemMessage, HumanMessage, BaseMessage
from langgraph.graph import END, StateGraph
from langgraph.graph.state import CompiledStateGraph

try:
    import joblib
    import numpy as _np
    _LAYER2_AVAILABLE = True
except ImportError:
    _LAYER2_AVAILABLE = False

from prompts import (
    CLASSIFY_PROMPT_TEMPLATE,
    CODE_GENERATION_PROMPT,
    CONFIDENCE_PROMPT_TEMPLATE,
    CONVERSATIONAL_PROMPT,
    GENERATE_PROMPT,
    PLATFORM_PROMPT,
    REFORMULATE_PROMPT,
    TEST_GENERATION_PROMPT,
)

from state import AgentState, ClassifyEntry

logger = logging.getLogger(__name__)

# ── AVAP Parser client — ADR-0009 (PTVL) ──────────────────────────────────────

_PARSER_URL     = os.getenv("AVAP_PARSER_URL", "http://45.77.193.144:8888")
_PARSER_TIMEOUT = int(os.getenv("AVAP_PARSER_TIMEOUT", "2"))
_CB_THRESHOLD   = int(os.getenv("PARSER_CB_THRESHOLD", "3"))
_CB_COOLDOWN    = int(os.getenv("PARSER_CB_COOLDOWN", "30"))


class _CircuitBreaker:
    CLOSED    = "CLOSED"
    OPEN      = "OPEN"
    HALF_OPEN = "HALF_OPEN"

    def __init__(self, threshold: int, cooldown: float):
        self._state     = self.CLOSED
        self._failures  = 0
        self._opened_at = 0.0
        self._threshold = threshold
        self._cooldown  = cooldown
        self._lock      = threading.Lock()

    def allow(self) -> bool:
        with self._lock:
            if self._state == self.CLOSED:
                return True
            if self._state == self.OPEN:
                if time.time() - self._opened_at >= self._cooldown:
                    self._state = self.HALF_OPEN
                    logger.info("[ptvl] circuit HALF-OPEN — probing parser availability")
                    return True
                return False
            # HALF_OPEN: allow the probe
            return True

    def success(self):
        with self._lock:
            if self._state != self.CLOSED:
                logger.info("[ptvl] circuit CLOSED — parser reachable")
            self._state    = self.CLOSED
            self._failures = 0

    def failure(self):
        with self._lock:
            self._failures += 1
            if self._failures >= self._threshold or self._state == self.HALF_OPEN:
                self._state     = self.OPEN
                self._opened_at = time.time()
                logger.warning("[ptvl] circuit OPEN — skipping parser, returning unvalidated")


_parser_cb = _CircuitBreaker(_CB_THRESHOLD, _CB_COOLDOWN)


def _strip_thinking(text: str) -> str:
    """Remove qwen3 <think>...</think> blocks from LLM output."""
    text = _re.sub(r"<think>.*?</think>", "", text, flags=_re.DOTALL)
    # Also strip a lone closing tag if the model omitted the opening one
    text = _re.sub(r"</think>", "", text)
    return text.strip()


def _extract_avap_code(text: str) -> str:
    """Return the first AVAP code block found in an LLM response."""
    text = _strip_thinking(text)

    for pattern in (r'```avap\s*\n(.*?)```', r'```\s*\n(.*?)```', r'```(.*?)```'):
        m = _re.search(pattern, text, _re.DOTALL)
        if m:
            return m.group(1).strip()

    return text


def _call_parser(text: str, test_inputs: dict = None, test_list: list = None) -> tuple:
    """Call AVAP Parser REST API.

    Tries /api/v1/execute first (executes the code, catches runtime errors).
    If test_inputs/test_list are provided, also validates assertions.
    Falls back to /parse (AST-only validation) if /api/v1/execute is not available.

    Returns:
        (True,  "")      — code valid and executed successfully (assertions passed if provided)
        (False, trace)   — code invalid, runtime error, or failed assertions
        (None,  "")      — parser unavailable or circuit open
    """

    if not _PARSER_URL or _PARSER_TIMEOUT == 0:
        return None, ""

    if not _parser_cb.allow():
        return None, ""

    code = _extract_avap_code(text)
    logger.info(f"[ptvl] extracted code ({len(code)} chars): {repr(code[:120])}")
    if not code.strip():
        return None, ""

    base_url = _PARSER_URL.rstrip('/')

    try:
        payload = {
            "code": code,
            "test_inputs": test_inputs or {},
            "test_list": test_list or [],
        }
        resp = _requests.post(
            f"{base_url}/api/v1/run",
            json=payload,
            timeout=_PARSER_TIMEOUT,
        )

        if resp.status_code == 404:
            # /api/v1/run not deployed yet — fall back to /parse
            logger.info("[ptvl] /api/v1/run not available, falling back to /parse")
            return _call_parser_parse(base_url, code)

        data = resp.json()
        _parser_cb.success()
        logger.info(f"[ptvl] parser response: {data}")

        if not data.get("success", False):
            error = data.get("error", "")
            if not error:
                failed_logs = [l for l in data.get("logs", []) if not l.get("success")]
                if failed_logs:
                    error = failed_logs[0].get("error", "runtime error")

            # Detect a KeyError from assertion evaluation — pattern: "'variable_name'"
            # This means the assertion referenced a variable that doesn't exist in scope
            # (e.g. str(result) when the code never assigns `result`).
            # The code itself is fine; drop the bad assertion and re-validate.
            _key_err_m = _re.match(r"^'([A-Za-z_]\w*)'$", error or "")
            if _key_err_m and test_list:
                bad_var  = _key_err_m.group(1)
                logger.warning(
                    f"[ptvl] assertion KeyError: variable '{bad_var}' not in scope — "
                    "dropping affected assertions (test defect, not code defect)"
                )
                filtered = [a for a in test_list if f"str({bad_var})" not in a]
                if filtered != test_list:
                    # Retry with the cleaned assertion list
                    return _call_parser(text, test_inputs=test_inputs, test_list=filtered)
                # All assertions referenced the bad variable — code execution was fine
                return True, ""

            return False, error or "runtime error"

        # Execution succeeded — check assertion result if assertions were provided
        if test_list and not data.get("assertion_result", True):
            return False, "assertion failed: the code ran but did not produce the expected output"

        return True, ""

    except Exception as exc:
        _parser_cb.failure()
        logger.warning(f"[ptvl] parser call failed: {exc}")
        return None, ""


def _call_parser_parse(base_url: str, code: str) -> tuple:
    """AST-only fallback via /parse."""
    try:
        resp = _requests.post(
            f"{base_url}/parse",
            json={"code": code},
            timeout=_PARSER_TIMEOUT,
        )
        data = resp.json()
        _parser_cb.success()
        logger.info(f"[ptvl] parser response (/parse fallback): {data}")
        if data.get("valid", False):
            return True, ""
        return False, data.get("error", "parse error")
    except Exception as exc:
        _parser_cb.failure()
        logger.warning(f"[ptvl] /parse fallback failed: {exc}")
        return None, ""


# ── Test generation helper — used by both build_graph and AskAgentStream ───────

def _run_generate_tests(user_request: str, generated_code: str, llm) -> tuple:
    """Generate test_inputs + test_list for the given AVAP code.

    Returns (test_inputs: dict, test_list: list).  Never raises — falls back to
    ({}, []) on timeout or any LLM/parse error so the caller can still validate.
    """
    import concurrent.futures

    def _run():
        prompt = "/no_think\n\n" + TEST_GENERATION_PROMPT.format(
            user_request=user_request,
            generated_code=generated_code,
        )
        resp = llm.invoke([SystemMessage(content=prompt)])
        raw  = _strip_thinking(resp.content)
        logger.info(f"[generate_tests] raw output: {repr(raw[:200])}")

        if raw.startswith("```"):
            raw = _re.sub(r"^```[a-z]*\n?", "", raw)
            raw = raw.rstrip("`").strip()

        m = _re.search(r'\{.*\}', raw, _re.DOTALL)
        if m:
            raw = m.group(0)

        data = _json.loads(raw)
        if not isinstance(data, dict):
            raise ValueError(f"expected JSON object, got {type(data).__name__}: {repr(raw[:80])}")
        return data.get("test_inputs", {}), data.get("test_list", [])

    try:
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex:
            future = ex.submit(_run)
            test_inputs, test_list = future.result(timeout=15)
        logger.info(f"[generate_tests] {len(test_list)} assertions generated")
    except concurrent.futures.TimeoutError:
        logger.warning("[generate_tests] timed out after 15s — skipping assertions")
        test_inputs, test_list = {}, []
    except Exception as exc:
        logger.warning(f"[generate_tests] skipped ({type(exc).__name__}): {exc}")
        test_inputs, test_list = {}, []

    return test_inputs, test_list


# ── Session stores ─────────────────────────────────────────────────────────────

session_store: dict[str, list] = defaultdict(list)
classify_history_store: dict[str, list] = defaultdict(list)


def format_context(docs):
    chunks = []
    for i, doc in enumerate(docs, 1):
        meta = doc.metadata or {}
        chunk_id = meta.get("chunk_id", meta.get("id", f"chunk-{i}"))
        source = meta.get("source_file", meta.get("source", "unknown"))
        doc_type = meta.get("doc_type", "")
        block_type = meta.get("block_type", "")
        section = meta.get("section", "")

        text = (doc.page_content or "").strip()
        if not text:
            text = meta.get("content") or meta.get("text") or ""

        header_parts = [f"[{i}]", f"id={chunk_id}"]
        if doc_type: header_parts.append(f"type={doc_type}")
        if block_type: header_parts.append(f"block={block_type}")
        if section: header_parts.append(f"section={section}")
        header_parts.append(f"source={source}")

        if doc_type in ("code", "code_example", "bnf") or \
           block_type in ("function", "if", "startLoop", "try"):
            header_parts.append("[AVAP CODE]")

        chunks.append(" ".join(header_parts) + "\n" + text)

    return "\n\n".join(chunks)


def format_history_for_classify(messages):
    lines = []
    for msg in messages[-6:]:
        if isinstance(msg, HumanMessage):
            lines.append(f"User: {msg.content}")
        elif isinstance(msg, AIMessage):
            lines.append(f"Assistant: {msg.content[:300]}")
        elif isinstance(msg, dict):
            role = msg.get("role", "user")
            content = msg.get("content", "")[:300]
            lines.append(f"{role.capitalize()}: {content}")
    return "\n".join(lines) if lines else "(no history)"


def hybrid_search_native(es_client, embeddings, query, index_name, k=8):
    query_vector = None
    try:
        query_vector = embeddings.embed_query(query)
    except Exception as e:
        logger.warning(f"[hybrid] embed_query fails: {e}")

    bm25_hits = []
    try:
        resp = es_client.search(
            index=index_name,
            body={
                "size": k,
                "query": {
                    "multi_match": {
                        "query": query,
                        "fields": ["content^2", "text^2"],
                        "type": "best_fields",
                        "fuzziness": "AUTO",
                    }
                },
                "_source": {"excludes": ["embedding"]},
            }
        )
        bm25_hits = resp["hits"]["hits"]
        logger.info(f"[hybrid] BM25 -> {len(bm25_hits)} hits")
    except Exception as e:
        logger.warning(f"[hybrid] BM25 fails: {e}")

    knn_hits = []
    if query_vector:
        try:
            resp = es_client.search(
                index=index_name,
                body={
                    "size": k,
                    "knn": {
                        "field": "embedding",
                        "query_vector": query_vector,
                        "k": k,
                        "num_candidates": k * 5,
                    },
                    "_source": {"excludes": ["embedding"]},
                }
            )
            knn_hits = resp["hits"]["hits"]
            logger.info(f"[hybrid] kNN -> {len(knn_hits)} hits")
        except Exception as e:
            logger.warning(f"[hybrid] kNN fails: {e}")

    rrf_scores: dict[str, float] = defaultdict(float)
    hit_by_id:  dict[str, dict]  = {}

    for rank, hit in enumerate(bm25_hits):
        doc_id = hit["_id"]
        rrf_scores[doc_id] += 1.0 / (rank + 60)
        hit_by_id[doc_id] = hit

    for rank, hit in enumerate(knn_hits):
        doc_id = hit["_id"]
        rrf_scores[doc_id] += 1.0 / (rank + 60)
        if doc_id not in hit_by_id:
            hit_by_id[doc_id] = hit

    ranked = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)[:k]

    docs = []
    for doc_id, score in ranked:
        src  = hit_by_id[doc_id]["_source"]
        text = src.get("content") or src.get("text") or ""
        meta = {k: v for k, v in src.items()
                if k not in ("content", "text", "embedding")}
        meta["id"]= doc_id
        meta["rrf_score"] = score
        docs.append(Document(page_content=text, metadata=meta))

    logger.info(f"[hybrid] RRF -> {len(docs)} final docs")
    return docs

def _build_classify_prompt(question: str, history_text: str, selected_text: str) -> str:
    prompt = (
        CLASSIFY_PROMPT_TEMPLATE
        .replace("{history}", history_text)
        .replace("{message}", question)
    )
    if selected_text:
        editor_section = (
            "\n\n<editor_selection>\n"
            "The user currently has the following AVAP code selected in their editor. "
            "If the question refers to 'this', 'here', 'the code above', or similar, "
            "it is about this selection.\n"
            f"{selected_text}\n"
            "</editor_selection>"
        )

        prompt = prompt.replace(
            f"<user_message>{question}</user_message>",
            f"{editor_section}\n\n<user_message>{question}</user_message>"
        )
    return prompt


def _build_reformulate_query(question: str, selected_text: str) -> str:
    if not selected_text:
        return question
    return f"{selected_text}\n\nUser question about the above: {question}"


def _build_generation_prompt(template_prompt: SystemMessage, context: str,
                              editor_content: str, selected_text: str,
                              extra_context: str) -> SystemMessage:
    base = template_prompt.content.format(context=context)

    sections = []

    if selected_text:
        sections.append(
            "<selected_code>\n"
            "The user has the following AVAP code selected in their editor. "
            "Ground your answer in this code first. "
            "Use the RAG context as supplementary reference only.\n"
            f"{selected_text}\n"
            "</selected_code>"
        )

    if editor_content:
        sections.append(
            "<editor_file>\n"
            "Full content of the active file open in the editor "
            "(use for broader context if needed):\n"
            f"{editor_content}\n"
            "</editor_file>"
        )

    if extra_context:
        sections.append(
            "<extra_context>\n"
            f"{extra_context}\n"
            "</extra_context>"
        )

    if sections:
        editor_block = "\n\n".join(sections)
        base = editor_block + "\n\n" + base

    return SystemMessage(content=base)


def _format_intent_history(classify_history: list) -> str:
    if not classify_history:
        return "(no prior turns)"
    lines = [f"[{e['type']}] \"{e['topic']}\"" for e in classify_history[-6:]]
    return "\n".join(lines)


_RC01_PREFIXES = (
    # System-injected prefixes — always PLATFORM (RC-01)
    "you are a direct and concise assistant",
)

_RC02_KEYWORDS = (
    # Account / subscription vocabulary (RC-02)
    "usage percentage",
    "project usage",
    "account usage",
    "your quota",
    "your subscription",
    "your plan",
    "your account",
    "your limits",
    "your consumption",
    "api calls",
    "llamadas este mes",
    "uso de la cuenta",
    "plan contratado",
    "límite de",
    "limite de",
)

# RC-02 structural signals: "N%" in context of usage/quota/consumption
_RC02_PERCENTAGE_RE = _re.compile(
    r'\b\d+(\.\d+)?\s*%',  # any percentage figure
)
_RC02_USAGE_CONTEXT_WORDS = (
    "usage", "quota", "consumption", "limit", "billing",
    "uso", "cuota", "consumo", "límite", "facturación",
    "project", "account", "plan", "subscription",
)


def _is_platform_query(question: str) -> bool:
    """Deterministic fast-path for PLATFORM classification.

    RC-01: known system-injected prefixes → always PLATFORM, no LLM call.
    RC-02: messages containing account/metrics/usage signals → always PLATFORM,
           regardless of conversation history or LLM classifier output.

    Both rules are hard-coded here, not in the prompt, so they cannot be
    overridden by model behavior.
    """
    q = question.strip().lower()

    # RC-01 — known platform-injected prefixes
    if any(prefix in q for prefix in _RC01_PREFIXES):
        return True

    # RC-02 — vocabulary signals
    if any(kw in q for kw in _RC02_KEYWORDS):
        return True

    # RC-02 — structural signal: percentage figure + usage context word
    if _RC02_PERCENTAGE_RE.search(q):
        if any(ctx in q for ctx in _RC02_USAGE_CONTEXT_WORDS):
            return True

    return False


# ── Layer 2 — embedding classifier (ADR-0008 Phase 2) ─────────────────────────

_LAYER2_MODEL_PATH          = os.getenv("CLASSIFIER_MODEL_PATH", "/data/classifier_model.pkl")
_LAYER2_CONFIDENCE_THRESHOLD = float(os.getenv("CLASSIFIER_CONFIDENCE_THRESHOLD", "0.85"))
_layer2_model: dict | None  = None


def _load_layer2_model() -> None:
    """Load the serialized Layer 2 classifier from disk (called once at startup).

    Safe to call even if the model file does not exist yet — the engine falls
    back to the LLM classifier (Layer 3) for every request.
    """
    global _layer2_model
    if not _LAYER2_AVAILABLE:
        logger.info("[classifier/L2] joblib/numpy not available — Layer 2 disabled")
        return
    p = Path(_LAYER2_MODEL_PATH)
    if p.exists():
        _layer2_model = joblib.load(p)
        meta = {k: v for k, v in _layer2_model.items() if k != "clf"}
        logger.info(f"[classifier/L2] model loaded from {p} — {meta}")
    else:
        logger.info(f"[classifier/L2] model not found at {p} — using LLM fallback only")


def _classify_layer2(query: str, embeddings) -> str | None:
    """Embed query and classify via Layer 2 model.

    Returns the predicted type if confidence >= threshold, else None (falls
    through to the LLM classifier, Layer 3).

    Never raises — any error returns None to preserve the Layer 3 fallback.
    """
    if _layer2_model is None or not _LAYER2_AVAILABLE:
        return None
    try:
        vec   = _np.array(embeddings.embed_query(query)).reshape(1, -1)
        proba = _layer2_model["clf"].predict_proba(vec)[0]
        conf  = float(proba.max())
        if conf >= _LAYER2_CONFIDENCE_THRESHOLD:
            label = _layer2_model["classes"][int(proba.argmax())]
            logger.info(f"[classifier/L2] → {label} (conf={conf:.2f})")
            return label
        logger.info(f"[classifier/L2] low confidence ({conf:.2f}) → LLM fallback")
        return None
    except Exception as e:
        logger.warning(f"[classifier/L2] error: {e} — falling through to LLM")
        return None


def _parse_query_type(raw: str) -> tuple[str, bool]:
    parts = raw.strip().upper().split()
    query_type = "RETRIEVAL"
    use_editor = False
    if parts:
        first = parts[0]
        if first.startswith("CODE_GENERATION") or "CODE" in first:
            query_type = "CODE_GENERATION"
        elif first.startswith("CONVERSATIONAL"):
            query_type = "CONVERSATIONAL"
        elif first.startswith("PLATFORM"):
            query_type = "PLATFORM"
    if len(parts) > 1 and parts[1] == "EDITOR":
        use_editor = True
    return query_type, use_editor

def build_graph(llm, embeddings, es_client, index_name, llm_conversational=None):

    _llm_conv = llm_conversational or llm

    def _persist(state: AgentState, response: BaseMessage):
        session_id = state.get("session_id", "")
        if session_id:
            session_store[session_id] = list(state["messages"]) + [response]

    def classify(state):
        messages = state["messages"]
        user_msg = messages[-1]
        question = getattr(user_msg, "content",
                           user_msg.get("content", "")
                           if isinstance(user_msg, dict) else "")
        selected_text    = state.get("selected_text", "")
        classify_history = state.get("classify_history") or []

        topic_snippet = question.strip()[:60].replace("\n", " ")

        # Phase 3 — caller-declared type: bypass all layers (ADR-0008 RC)
        declared = state.get("query_type", "")
        if declared in ("RETRIEVAL", "CODE_GENERATION", "CONVERSATIONAL", "PLATFORM"):
            logger.info(f"[classify] caller-declared → {declared}")
            entry: ClassifyEntry = {"type": declared, "topic": topic_snippet}
            return {
                "query_type": declared,
                "use_editor_context": False,
                "classify_history": classify_history + [entry],
            }

        # Layer 1 — hard rules (RC-01, RC-02)
        if _is_platform_query(question):
            logger.info("[classify] L1 → PLATFORM")
            entry: ClassifyEntry = {"type": "PLATFORM", "topic": topic_snippet}
            return {
                "query_type": "PLATFORM",
                "use_editor_context": False,
                "classify_history": classify_history + [entry],
            }

        # Layer 2 — embedding classifier (ADR-0008 Phase 2)
        l2_type = _classify_layer2(question, embeddings)
        if l2_type:
            entry: ClassifyEntry = {"type": l2_type, "topic": topic_snippet}
            return {
                "query_type": l2_type,
                "use_editor_context": False,
                "classify_history": classify_history + [entry],
            }

        # Layer 3 — LLM classifier (bootstrap fallback)
        intent_history_text = _format_intent_history(classify_history)
        prompt_content = _build_classify_prompt(question, intent_history_text, selected_text)

        resp = llm.invoke([SystemMessage(content=prompt_content)])
        raw  = resp.content.strip().upper()
        query_type, use_editor_ctx = _parse_query_type(raw)
        logger.info(f"[classify] L3 selected={bool(selected_text)} raw='{raw}' -> {query_type} editor={use_editor_ctx}")

        entry: ClassifyEntry = {"type": query_type, "topic": topic_snippet}
        return {
            "query_type": query_type,
            "use_editor_context": use_editor_ctx,
            "classify_history": classify_history + [entry],
        }

    def reformulate(state: AgentState) -> AgentState:
        user_msg = state["messages"][-1]
        selected_text = state.get("selected_text", "")
        question = getattr(user_msg, "content",
                                user_msg.get("content", "")
                                if isinstance(user_msg, dict) else "")

        anchor = _build_reformulate_query(question, selected_text)

        if selected_text:

            from langchain_core.messages import HumanMessage as HM
            resp = llm.invoke([REFORMULATE_PROMPT, HM(content=anchor)])
        else:
            query_type = state.get("query_type", "RETRIEVAL")
            mode_hint = HumanMessage(content=f"[MODE: {query_type}]\n{question}")
            resp = llm.invoke([REFORMULATE_PROMPT, mode_hint])

        reformulated = resp.content.strip()
        logger.info(f"[reformulate] selected={bool(selected_text)} -> '{reformulated}'")
        return {"reformulated_query": reformulated}

    def retrieve(state: AgentState) -> AgentState:
        query = state["reformulated_query"]
        docs = hybrid_search_native(
            es_client=es_client,
            embeddings=embeddings,
            query=query,
            index_name=index_name,
            k=8,
        )
        context = format_context(docs)
        logger.info(f"[retrieve] {len(docs)} docs, context len={len(context)}")
        return {"context": context}

    def generate(state):
        use_editor = state.get("use_editor_context", False)
        prompt = _build_generation_prompt(
            template_prompt = GENERATE_PROMPT,
            context = state.get("context", ""),
            editor_content  = state.get("editor_content", "") if use_editor else "",
            selected_text = state.get("selected_text", "") if use_editor else "",
            extra_context = state.get("extra_context", ""),
        )
        resp = llm.invoke([prompt] + state["messages"])
        logger.info(f"[generate] {len(resp.content)} chars")
        _persist(state, resp)
        return {"messages": [resp]}

    def generate_code(state):
        use_editor = state.get("use_editor_context", False)
        prompt = _build_generation_prompt(
            template_prompt = CODE_GENERATION_PROMPT,
            context = state.get("context", ""),
            editor_content = state.get("editor_content", "") if use_editor else "",
            selected_text = state.get("selected_text", "") if use_editor else "",
            extra_context = state.get("extra_context", ""),
        )
        resp = llm.invoke([prompt] + state["messages"])
        logger.info(f"[generate_code] {len(resp.content)} chars")
        #logger.info(resp.content)
        _persist(state, resp)
        return {"messages": [resp]}

    def generate_tests(state: AgentState) -> AgentState:
        user_msg       = state["messages"][-2] if len(state["messages"]) >= 2 else state["messages"][-1]
        generated      = state["messages"][-1]
        user_request   = getattr(user_msg,  "content", str(user_msg))
        generated_code = getattr(generated, "content", str(generated))
        test_inputs, test_list = _run_generate_tests(user_request, generated_code, llm)
        return {"test_inputs": test_inputs, "test_list": test_list}

    def respond_conversational(state):
        extra_context = state.get("extra_context", "")
        if extra_context:
            enriched = SystemMessage(content=(
                CONVERSATIONAL_PROMPT.content +
                f"\n\n<extra_context>\n{extra_context}\n</extra_context>"
            ))
        else:
            enriched = CONVERSATIONAL_PROMPT
        resp = _llm_conv.invoke([enriched] + state["messages"])
        logger.info("[conversational] from conversation")
        _persist(state, resp)
        return {"messages": [resp]}

    def respond_platform(state):
        extra_context = state.get("extra_context", "")
        if extra_context:
            enriched = SystemMessage(content=(
                PLATFORM_PROMPT.content +
                f"\n\n<extra_context>\n{extra_context}\n</extra_context>"
            ))
        else:
            enriched = PLATFORM_PROMPT
        resp = _llm_conv.invoke([enriched] + state["messages"])
        logger.info("[platform] platform/account query")
        _persist(state, resp)
        return {"messages": [resp]}

    # ── PTVL nodes (ADR-0009) ─────────────────────────────────────────────────

    def validate_code(state: AgentState) -> AgentState:
        last_msg = state["messages"][-1]
        content  = getattr(last_msg, "content", str(last_msg))
        valid, trace = _call_parser(
            content,
            test_inputs=state.get("test_inputs") or {},
            test_list=state.get("test_list") or [],
        )
        if valid is None:
            logger.warning("[ptvl] parser unavailable — returning unvalidated")
            return {"validation_status": "PARSER_UNAVAILABLE", "parser_trace": ""}
        if valid:
            logger.info("[ptvl] code VALID on first attempt")
            return {"validation_status": "", "parser_trace": ""}
        logger.info(f"[ptvl] code INVALID — trace: {str(trace)[:120]}")
        return {"validation_status": "INVALID", "parser_trace": str(trace)}

    def generate_code_retry(state: AgentState) -> AgentState:
        parser_trace = state.get("parser_trace", "")
        use_editor   = state.get("use_editor_context", False)

        feedback = (
            "\n\n<parser_feedback>\n"
            "The previous attempt failed when the AVAP code was executed. "
            "The execution engine reported the following error:\n\n"
            f"{parser_trace}\n\n"
            "Rules to fix it:\n"
            "- Read the error carefully — it identifies the exact command or line that failed.\n"
            "- Do NOT repeat the same construct that caused the error.\n"
            "- Only use commands from <avap_syntax_reminder> or <context>.\n"
            "- If the error mentions an unknown command, replace it with the correct AVAP equivalent.\n"
            "- If the error mentions a variable, make sure it is declared before use.\n"
            "</parser_feedback>"
        ) if parser_trace else ""

        base = _build_generation_prompt(
            template_prompt=CODE_GENERATION_PROMPT,
            context=state.get("context", ""),
            editor_content=state.get("editor_content", "") if use_editor else "",
            selected_text=state.get("selected_text", "")   if use_editor else "",
            extra_context=state.get("extra_context", ""),
        )
        retry_prompt = SystemMessage(content=base.content + feedback)
        resp = llm.invoke([retry_prompt] + state["messages"])
        logger.info(f"[generate_code_retry] {len(resp.content)} chars")
        _persist(state, resp)
        return {"messages": [resp]}

    def validate_code_after_retry(state: AgentState) -> AgentState:
        last_msg = state["messages"][-1]
        content  = getattr(last_msg, "content", str(last_msg))
        valid, trace = _call_parser(content)
        if valid is None:
            logger.warning("[ptvl] parser unavailable after retry")
            return {"validation_status": "PARSER_UNAVAILABLE", "parser_trace": ""}
        if valid:
            logger.info("[ptvl] retry code VALID")
            return {"validation_status": "", "parser_trace": ""}
        logger.info("[ptvl] retry code still INVALID → INVALID_UNRESOLVED")
        return {"validation_status": "INVALID_UNRESOLVED", "parser_trace": str(trace)}

    def check_context_relevance(state: AgentState) -> AgentState:
        user_msg = state["messages"][-1]
        question = getattr(user_msg, "content",
                           user_msg.get("content", "") if isinstance(user_msg, dict) else "")
        context  = state.get("context", "")
        prompt   = CONFIDENCE_PROMPT_TEMPLATE.format(question=question, context=context)
        resp     = llm.invoke([SystemMessage(content=prompt)])
        relevant = resp.content.strip().upper().startswith("YES")
        logger.info(f"[check_context_relevance] relevant={relevant}")
        return {"context_relevant": relevant}

    def reformulate_with_hint(state: AgentState) -> AgentState:
        user_msg = state["messages"][-1]
        question = getattr(user_msg, "content",
                           user_msg.get("content", "") if isinstance(user_msg, dict) else "")
        hint = (
            f"[CONTEXT_INSUFFICIENT]\n"
            f"The previous retrieval did not return relevant context for: \"{question}\"\n"
            f"Reformulate this query using broader terms or alternative phrasing."
        )
        resp = llm.invoke([REFORMULATE_PROMPT, HumanMessage(content=hint)])
        reformulated = resp.content.strip()
        logger.info(f"[reformulate_with_hint] -> '{reformulated}'")
        return {"reformulated_query": reformulated}

    def retrieve_retry(state: AgentState) -> AgentState:
        query = state["reformulated_query"]
        docs  = hybrid_search_native(
            es_client=es_client, embeddings=embeddings,
            query=query, index_name=index_name, k=8,
        )
        context = format_context(docs)
        logger.info(f"[retrieve_retry] {len(docs)} docs, context len={len(context)}")
        return {"context": context}

    # ── Routing ───────────────────────────────────────────────────────────────

    def route_by_type(state):
        return state.get("query_type", "RETRIEVAL")

    def route_after_retrieve(state):
        qt = state.get("query_type", "RETRIEVAL")
        return "generate_code" if qt == "CODE_GENERATION" else "check_context_relevance"

    def route_after_validate(state):
        if state.get("validation_status", "") == "INVALID":
            return "generate_code_retry"
        return END

    def route_after_context_check(state):
        if state.get("context_relevant", True):
            return "generate"
        return "reformulate_with_hint"

    graph_builder = StateGraph(AgentState)

    graph_builder.add_node("classify",                classify)
    graph_builder.add_node("reformulate",              reformulate)
    graph_builder.add_node("retrieve",                 retrieve)
    graph_builder.add_node("generate",                 generate)
    graph_builder.add_node("generate_code",            generate_code)
    graph_builder.add_node("generate_tests",           generate_tests)
    graph_builder.add_node("validate_code",            validate_code)
    graph_builder.add_node("generate_code_retry",      generate_code_retry)
    graph_builder.add_node("validate_code_after_retry",validate_code_after_retry)
    graph_builder.add_node("check_context_relevance",  check_context_relevance)
    graph_builder.add_node("reformulate_with_hint",    reformulate_with_hint)
    graph_builder.add_node("retrieve_retry",           retrieve_retry)
    graph_builder.add_node("respond_conversational",   respond_conversational)
    graph_builder.add_node("respond_platform",         respond_platform)

    graph_builder.set_entry_point("classify")

    graph_builder.add_conditional_edges(
        "classify",
        route_by_type,
        {
            "RETRIEVAL":       "reformulate",
            "CODE_GENERATION": "reformulate",
            "CONVERSATIONAL":  "respond_conversational",
            "PLATFORM":        "respond_platform",
        }
    )

    graph_builder.add_edge("reformulate", "retrieve")

    graph_builder.add_conditional_edges(
        "retrieve",
        route_after_retrieve,
        {
            "generate_code":           "generate_code",
            "check_context_relevance": "check_context_relevance",
        }
    )

    # CODE_GENERATION path: generate → generate_tests → validate → (retry if invalid) → END
    graph_builder.add_edge("generate_code",   "generate_tests")
    graph_builder.add_edge("generate_tests",  "validate_code")
    graph_builder.add_conditional_edges(
        "validate_code",
        route_after_validate,
        {
            "generate_code_retry": "generate_code_retry",
            END:                   END,
        }
    )
    graph_builder.add_edge("generate_code_retry", "validate_code_after_retry")
    graph_builder.add_edge("validate_code_after_retry", END)

    # RETRIEVAL path: context check → generate or reformulate+retry → generate
    graph_builder.add_conditional_edges(
        "check_context_relevance",
        route_after_context_check,
        {
            "generate":              "generate",
            "reformulate_with_hint": "reformulate_with_hint",
        }
    )
    graph_builder.add_edge("reformulate_with_hint", "retrieve_retry")
    graph_builder.add_edge("retrieve_retry", "generate")

    graph_builder.add_edge("generate", END)
    graph_builder.add_edge("respond_conversational", END)
    graph_builder.add_edge("respond_platform", END)

    return graph_builder.compile()


def build_prepare_graph(llm, embeddings, es_client, index_name):

    def classify(state):
        messages = state["messages"]
        user_msg = messages[-1]
        question = getattr(user_msg, "content",
                                user_msg.get("content", "")
                                if isinstance(user_msg, dict) else "")
        history_msgs  = messages[:-1]
        selected_text = state.get("selected_text", "")

        classify_history = state.get("classify_history") or []
        topic_snippet    = question.strip()[:60].replace("\n", " ")

        # Phase 3 — caller-declared type: bypass all layers (ADR-0008 RC)
        declared = state.get("query_type", "")
        if declared in ("RETRIEVAL", "CODE_GENERATION", "CONVERSATIONAL", "PLATFORM"):
            logger.info(f"[prepare/classify] caller-declared → {declared}")
            entry: ClassifyEntry = {"type": declared, "topic": topic_snippet}
            return {
                "query_type": declared,
                "use_editor_context": False,
                "classify_history": classify_history + [entry],
            }

        # Layer 1 — hard rules (RC-01, RC-02)
        if _is_platform_query(question):
            logger.info("[prepare/classify] L1 → PLATFORM")
            entry: ClassifyEntry = {"type": "PLATFORM", "topic": topic_snippet}
            return {
                "query_type": "PLATFORM",
                "use_editor_context": False,
                "classify_history": classify_history + [entry],
            }

        # Layer 2 — embedding classifier (ADR-0008 Phase 2)
        l2_type = _classify_layer2(question, embeddings)
        if l2_type:
            entry: ClassifyEntry = {"type": l2_type, "topic": topic_snippet}
            return {
                "query_type": l2_type,
                "use_editor_context": False,
                "classify_history": classify_history + [entry],
            }

        # Layer 3 — LLM classifier (bootstrap fallback)
        intent_history_text = _format_intent_history(classify_history)
        prompt_content      = _build_classify_prompt(question, intent_history_text, selected_text)

        resp = llm.invoke([SystemMessage(content=prompt_content)])
        raw  = resp.content.strip().upper()
        query_type, use_editor_ctx = _parse_query_type(raw)
        logger.info(f"[prepare/classify] L3 selected={bool(selected_text)} raw='{raw}' -> {query_type} editor={use_editor_ctx}")
        entry: ClassifyEntry = {"type": query_type, "topic": topic_snippet}
        return {
            "query_type": query_type,
            "use_editor_context": use_editor_ctx,
            "classify_history": classify_history + [entry],
        }

    def reformulate(state: AgentState) -> AgentState:
        user_msg = state["messages"][-1]
        selected_text = state.get("selected_text", "")
        question = getattr(user_msg, "content",
                                user_msg.get("content", "")
                                if isinstance(user_msg, dict) else "")

        anchor = _build_reformulate_query(question, selected_text)

        if selected_text:
            from langchain_core.messages import HumanMessage as HM
            resp = llm.invoke([REFORMULATE_PROMPT, HM(content=anchor)])
        else:
            query_type = state.get("query_type", "RETRIEVAL")
            mode_hint = HumanMessage(content=f"[MODE: {query_type}]\n{question}")
            resp = llm.invoke([REFORMULATE_PROMPT, mode_hint])

        reformulated = resp.content.strip()
        logger.info(f"[prepare/reformulate] selected={bool(selected_text)} -> '{reformulated}'")
        return {"reformulated_query": reformulated}

    def retrieve(state: AgentState) -> AgentState:
        query = state["reformulated_query"]
        docs = hybrid_search_native(
            es_client=es_client,
            embeddings=embeddings,
            query=query,
            index_name=index_name,
            k=8,
        )
        context = format_context(docs)
        logger.info(f"[prepare/retrieve] {len(docs)} docs, context len={len(context)}")
        return {"context": context}

    def skip_retrieve(state: AgentState) -> AgentState:
        return {"context": ""}

    def check_context_relevance(state: AgentState) -> AgentState:
        user_msg = state["messages"][-1]
        question = getattr(user_msg, "content",
                           user_msg.get("content", "") if isinstance(user_msg, dict) else "")
        context  = state.get("context", "")
        prompt   = CONFIDENCE_PROMPT_TEMPLATE.format(question=question, context=context)
        resp     = llm.invoke([SystemMessage(content=prompt)])
        relevant = resp.content.strip().upper().startswith("YES")
        logger.info(f"[prepare/check_context_relevance] relevant={relevant}")
        return {"context_relevant": relevant}

    def reformulate_with_hint(state: AgentState) -> AgentState:
        user_msg = state["messages"][-1]
        question = getattr(user_msg, "content",
                           user_msg.get("content", "") if isinstance(user_msg, dict) else "")
        hint = (
            f"[CONTEXT_INSUFFICIENT]\n"
            f"The previous retrieval did not return relevant context for: \"{question}\"\n"
            f"Reformulate this query using broader terms or alternative phrasing."
        )
        resp = llm.invoke([REFORMULATE_PROMPT, HumanMessage(content=hint)])
        reformulated = resp.content.strip()
        logger.info(f"[prepare/reformulate_with_hint] -> '{reformulated}'")
        return {"reformulated_query": reformulated}

    def retrieve_retry(state: AgentState) -> AgentState:
        query = state["reformulated_query"]
        docs  = hybrid_search_native(
            es_client=es_client, embeddings=embeddings,
            query=query, index_name=index_name, k=8,
        )
        context = format_context(docs)
        logger.info(f"[prepare/retrieve_retry] {len(docs)} docs, context len={len(context)}")
        return {"context": context}

    def route_by_type(state):
        return state.get("query_type", "RETRIEVAL")

    def route_after_retrieve(state):
        qt = state.get("query_type", "RETRIEVAL")
        if qt == "RETRIEVAL":
            return "check_context_relevance"
        return END  # CODE_GENERATION goes straight to END

    def route_after_context_check(state):
        if state.get("context_relevant", True):
            return END
        return "reformulate_with_hint"

    graph_builder = StateGraph(AgentState)

    graph_builder.add_node("classify",               classify)
    graph_builder.add_node("reformulate",             reformulate)
    graph_builder.add_node("retrieve",                retrieve)
    graph_builder.add_node("skip_retrieve",           skip_retrieve)
    graph_builder.add_node("check_context_relevance", check_context_relevance)
    graph_builder.add_node("reformulate_with_hint",   reformulate_with_hint)
    graph_builder.add_node("retrieve_retry",          retrieve_retry)

    graph_builder.set_entry_point("classify")

    graph_builder.add_conditional_edges(
        "classify",
        route_by_type,
        {
            "RETRIEVAL":       "reformulate",
            "CODE_GENERATION": "reformulate",
            "CONVERSATIONAL":  "skip_retrieve",
            "PLATFORM":        "skip_retrieve",
        }
    )

    graph_builder.add_edge("reformulate", "retrieve")

    graph_builder.add_conditional_edges(
        "retrieve",
        route_after_retrieve,
        {
            "check_context_relevance": "check_context_relevance",
            END:                       END,
        }
    )

    graph_builder.add_conditional_edges(
        "check_context_relevance",
        route_after_context_check,
        {
            END:                   END,
            "reformulate_with_hint": "reformulate_with_hint",
        }
    )
    graph_builder.add_edge("reformulate_with_hint", "retrieve_retry")
    graph_builder.add_edge("retrieve_retry", END)

    graph_builder.add_edge("skip_retrieve", END)

    return graph_builder.compile()


def build_final_messages(state: AgentState) -> list:
    query_type = state.get("query_type", "RETRIEVAL")
    context = state.get("context", "")
    messages = state.get("messages", [])
    editor_content = state.get("editor_content", "")
    selected_text = state.get("selected_text", "")
    extra_context = state.get("extra_context", "")

    if query_type == "CONVERSATIONAL":
        extra_context = state.get("extra_context", "")
        if extra_context:
            enriched = SystemMessage(content=(
                CONVERSATIONAL_PROMPT.content +
                f"\n\n<extra_context>\n{extra_context}\n</extra_context>"
            ))
        else:
            enriched = CONVERSATIONAL_PROMPT
        return [enriched] + messages

    if query_type == "PLATFORM":
        extra_context = state.get("extra_context", "")
        if extra_context:
            enriched = SystemMessage(content=(
                PLATFORM_PROMPT.content +
                f"\n\n<extra_context>\n{extra_context}\n</extra_context>"
            ))
        else:
            enriched = PLATFORM_PROMPT
        return [enriched] + messages

    use_editor = state.get("use_editor_context", False)
    if query_type == "CODE_GENERATION":
        prompt = _build_generation_prompt(
            template_prompt = CODE_GENERATION_PROMPT,
            context = context,
            editor_content = editor_content if use_editor else "",
            selected_text = selected_text  if use_editor else "",
            extra_context = extra_context,
        )
    else:
        prompt = _build_generation_prompt(
            template_prompt = GENERATE_PROMPT,
            context = context,
            editor_content  = editor_content if use_editor else "",
            selected_text= selected_text  if use_editor else "",
            extra_context = extra_context,
        )

    return [prompt] + messages
