"""LangGraph-based PRAL orchestrator and core nodes for the chatbot flow."""

from __future__ import annotations

import json
import logging
import re
import time
from collections.abc import AsyncGenerator
from typing import TYPE_CHECKING, Any, Dict, List, Literal, NotRequired, Protocol, TypedDict

from pydantic import BaseModel, Field

try:
    from langgraph.types import StreamWriter
except ImportError:
    StreamWriter = Any

try:
    from langgraph.graph import END, START, StateGraph
except ImportError:  # pragma: no cover - handled by runtime validation.
    END = None
    START = None
    StateGraph = None

from app.services.chatbot.chat_history import ChatHistoryManager

if TYPE_CHECKING:
    from app.services.chatbot.graph_retriever import GraphRetriever
    from app.services.chatbot.vector_retriever import VectorRetriever
    from app.services.chatbot.research_policy import ResearchPolicy

logger = logging.getLogger(__name__)


class ReferenceItem(TypedDict):
    """Structured reference payload for UI and persistence."""

    title: str
    url: str
    snippet: str


class RetrievedChunk(TypedDict):
    """Normalized chunk payload consumed across retrieval nodes."""

    file_hash: str
    title: str
    content: str
    metadata_header: NotRequired[str]
    score: NotRequired[float]
    source: NotRequired[str]


class NeighborDocument(TypedDict):
    """Neighbor document discovered from graph traversal."""

    hash: str
    title: str
    reasoning: str


class RoutingDecision(TypedDict):
    """Structured router output for deterministic branching."""

    layer: Literal[1, 2, 3]
    intent: str


class PRALState(TypedDict):
    """Shared state across Router, Retriever, Compressor, Generator, and Learn."""

    query: str
    project_id: str
    user_id: str
    active_hashes: List[str]
    mode: Literal["STRICT", "RESEARCH"]
    layer: Literal[1, 2, 3]
    intent: str
    target_hashes: List[str]
    neighbor_documents: List[NeighborDocument]
    graph_context_text: str
    raw_chunks: List[RetrievedChunk]
    compressed_chunks: List[RetrievedChunk]
    references: List[ReferenceItem]
    final_response: str
    metrics: Dict[str, Any]
    error_message: NotRequired[str]


class StructuredLLM(Protocol):
    """Protocol for structured-output capable chat models."""

    def with_structured_output(self, schema: type[BaseModel]) -> "StructuredInvoker":
        """Return an invoker that emits validated schema instances."""


class StructuredInvoker(Protocol):
    """Protocol for async structured model invokers."""

    async def ainvoke(self, payload: List[Dict[str, str]]) -> BaseModel:
        """Invoke the model asynchronously."""


class TextLLMResult(Protocol):
    """Protocol for text completion results."""

    @property
    def content(self) -> str:
        """Return generated text content."""


class TextLLM(Protocol):
    """Protocol for async text generation models."""

    async def ainvoke(self, payload: List[Dict[str, str]]) -> TextLLMResult:
        """Invoke the model asynchronously."""


class RouterOutput(BaseModel):
    """Validated router response."""

    layer: Literal[1, 2, 3] = Field(description="Selected retrieval layer.")
    intent: str = Field(description="Short explanation of the classified intent.")


class RouterNode:
    """LangGraph router node using structured LLM output."""

    def __init__(
        self,
        llm: StructuredLLM | None = None,
        model_name: str = "gpt-4o-mini",
    ) -> None:
        self._llm = llm or self._build_default_llm(model_name=model_name)

    @staticmethod
    def _build_default_llm(model_name: str) -> StructuredLLM:
        from app.services.ai.ai_gateway import create_ai_gateway

        gateway = create_ai_gateway()
        return gateway.get_langchain_model("chatbot", model_name=model_name)

    @staticmethod
    def _build_prompt(query: str) -> List[Dict[str, str]]:
        return [
            {
                "role": "system",
                "content": (
                    "Klasifikasikan pertanyaan ke salah satu layer retrieval.\n"
                    "Layer 1: fokus pada satu dokumen aktif.\n"
                    "Layer 2: perbandingan atau relasi antar dokumen.\n"
                    "Layer 3: pertanyaan umum di luar literatur aktif.\n"
                    "Kembalikan intent singkat dan objektif."
                ),
            },
            {"role": "user", "content": query},
        ]

    async def run(self, state: PRALState) -> PRALState:
        if state["mode"] == "STRICT":
            state["layer"] = 1
            state["intent"] = "STRICT mode forced Layer 1"
            state["metrics"]["layer_strategy"] = 1
            state["metrics"]["query_intent"] = state["intent"]
            return state

        structured_llm = self._llm.with_structured_output(RouterOutput)
        result = await structured_llm.ainvoke(self._build_prompt(state["query"]))

        state["layer"] = result.layer
        state["intent"] = result.intent
        state["metrics"]["layer_strategy"] = result.layer
        state["metrics"]["query_intent"] = result.intent
        return state


class RetrievalFlow:
    """LangGraph retrieval flow for graph expansion and vector chunk fetch."""

    def __init__(
        self,
        graph_retriever: "GraphRetriever | None" = None,
        vector_retriever: "VectorRetriever | None" = None,
        top_k: int = 15,
        max_neighbors: int = 3,
    ) -> None:
        if graph_retriever is None:
            from app.services.chatbot.graph_retriever import GraphRetriever

            graph_retriever = GraphRetriever()
        if vector_retriever is None:
            from app.services.chatbot.vector_retriever import VectorRetriever

            vector_retriever = VectorRetriever()
        self.graph_retriever = graph_retriever
        self.vector_retriever = vector_retriever
        self.top_k = top_k
        self.max_neighbors = max_neighbors

    @staticmethod
    def _normalize_chunks(raw_chunks: List[Dict[str, Any]]) -> List[RetrievedChunk]:
        normalized: List[RetrievedChunk] = []
        for chunk in raw_chunks:
            normalized.append(
                RetrievedChunk(
                    file_hash=chunk.get("file_hash", ""),
                    title=chunk.get("title", ""),
                    content=chunk.get("content", ""),
                    metadata_header=chunk.get("metadata_header", ""),
                    score=chunk.get("similarity") or chunk.get("score"),
                    source="supabase.match_document_chunks",
                )
            )
        return normalized

    async def run(self, state: PRALState) -> PRALState:
        target_hashes = list(state["active_hashes"])
        neighbor_documents: List[NeighborDocument] = []
        graph_context_lines: List[str] = []

        if state["layer"] == 2 and state["active_hashes"]:
            seen_hashes = set(target_hashes)
            for active_hash in state["active_hashes"]:
                neighbors = await self.graph_retriever.get_relevant_neighbors(
                    active_hash,
                    project_id=state["project_id"],
                    user_id=state["user_id"],
                    max_nodes=self.max_neighbors,
                )
                for neighbor in neighbors:
                    neighbor_doc = NeighborDocument(
                        hash=neighbor.get("hash", ""),
                        title=neighbor.get("title", ""),
                        reasoning=neighbor.get("reasoning", ""),
                    )
                    neighbor_documents.append(neighbor_doc)
                    if neighbor_doc["hash"] and neighbor_doc["hash"] not in seen_hashes:
                        target_hashes.append(neighbor_doc["hash"])
                        seen_hashes.add(neighbor_doc["hash"])
                    if neighbor_doc["title"] or neighbor_doc["reasoning"]:
                        graph_context_lines.append(
                            (
                                f"Terdapat koneksi dengan literatur "
                                f"'{neighbor_doc['title']}'. "
                                f"Alasan (Reasoning): {neighbor_doc['reasoning']}"
                            )
                        )

        raw_chunks: List[RetrievedChunk] = []
        if state["layer"] in (1, 2) and target_hashes:
            fetched_chunks = self.vector_retriever.similarity_search_multi_query(
                query=state["query"],
                hash_list=target_hashes,
                top_k=self.top_k,
            )
            raw_chunks = self._normalize_chunks(fetched_chunks)

        state["target_hashes"] = target_hashes
        state["neighbor_documents"] = neighbor_documents
        state["graph_context_text"] = (
            "\n".join(graph_context_lines) if graph_context_lines else ""
        )
        state["raw_chunks"] = raw_chunks
        state["metrics"]["source_count"] = len({chunk["file_hash"] for chunk in raw_chunks})
        return state


class CompressorNode:
    """Reduce noisy retrieval chunks using an injected or default mini model."""

    def __init__(
        self,
        llm: TextLLM | None = None,
        model_name: str = "gpt-4o-mini",
        max_chunks: int = 4,
    ) -> None:
        self._llm = llm or self._build_default_llm(model_name=model_name)
        self.max_chunks = max_chunks

    @staticmethod
    def _build_default_llm(model_name: str) -> TextLLM:
        from app.services.ai.ai_gateway import create_ai_gateway

        gateway = create_ai_gateway()
        return gateway.get_langchain_model("chatbot", model_name=model_name)

    @staticmethod
    def _chunk_digest(chunks: List[RetrievedChunk]) -> str:
        lines: List[str] = []
        for index, chunk in enumerate(chunks, start=1):
            header = chunk.get("metadata_header", "")
            lines.append(
                f"[{index}] hash={chunk['file_hash']} title={chunk['title']} "
                f"header={header}\n{chunk['content']}"
            )
        return "\n\n".join(lines)

    async def run(self, state: PRALState) -> PRALState:
        raw_chunks = state["raw_chunks"]
        if not raw_chunks:
            state["compressed_chunks"] = []
            state["metrics"]["compressed_chunk_count"] = 0
            return state

        prompt = [
            {
                "role": "system",
                "content": (
                    "Pilih potongan teks yang paling relevan untuk menjawab query.\n"
                    "Kembalikan hanya daftar indeks yang dipisahkan koma, maksimal "
                    f"{self.max_chunks} item."
                ),
            },
            {
                "role": "user",
                "content": (
                    f"Query: {state['query']}\n\n"
                    f"Chunks:\n{self._chunk_digest(raw_chunks)}"
                ),
            },
        ]
        result = await self._llm.ainvoke(prompt)
        selected_indexes: List[int] = []
        for token in result.content.replace("\n", ",").split(","):
            stripped = token.strip()
            if stripped.isdigit():
                idx = int(stripped)
                if 1 <= idx <= len(raw_chunks) and idx not in selected_indexes:
                    selected_indexes.append(idx)
        if not selected_indexes:
            selected_indexes = list(range(1, min(len(raw_chunks), self.max_chunks) + 1))

        compressed_chunks = [
            raw_chunks[index - 1] for index in selected_indexes[: self.max_chunks]
        ]
        state["compressed_chunks"] = compressed_chunks
        state["metrics"]["compressed_chunk_count"] = len(compressed_chunks)
        return state


class GeneratorNode:
    """Generate the final markdown response from compressed context."""

    def __init__(
        self,
        llm: TextLLM | None = None,
        policy: "ResearchPolicy | None" = None,
        graph_retriever: "GraphRetriever | None" = None,
        model_name: str = "gpt-4o-mini",
    ) -> None:
        self._llm = llm or self._build_default_llm(model_name=model_name)
        if policy is None:
            from app.services.chatbot.research_policy import ResearchPolicy

            policy = ResearchPolicy()
        if graph_retriever is None:
            from app.services.chatbot.graph_retriever import GraphRetriever

            graph_retriever = GraphRetriever()
        self.policy = policy
        self.graph_retriever = graph_retriever

    @staticmethod
    def _build_default_llm(model_name: str) -> TextLLM:
        from app.services.ai.ai_gateway import create_ai_gateway

        gateway = create_ai_gateway()
        return gateway.get_langchain_model("chatbot", model_name=model_name)

    @staticmethod
    def _is_placeholder_title(title: str | None) -> bool:
        if not title:
            return True
        normalized = title.strip().lower()
        return not normalized or normalized == "unknown" or normalized.endswith(".pdf")

    async def _resolve_reference_title(self, chunk: RetrievedChunk) -> str:
        file_hash = chunk["file_hash"]
        neo4j_title = await self.graph_retriever.get_node_title_by_hash(file_hash)
        if not self._is_placeholder_title(neo4j_title):
            return neo4j_title.strip()

        chunk_title = chunk.get("title", "")
        if not self._is_placeholder_title(chunk_title):
            return chunk_title.strip()

        return "Dokumen Akademis"

    async def _build_references(self, chunks: List[RetrievedChunk]) -> List[ReferenceItem]:
        seen_hashes: set[str] = set()
        resolved_titles: Dict[str, str] = {}
        references: List[ReferenceItem] = []
        for chunk in chunks:
            file_hash = chunk["file_hash"]
            if not file_hash or file_hash in seen_hashes:
                continue
            seen_hashes.add(file_hash)
            if file_hash not in resolved_titles:
                resolved_titles[file_hash] = await self._resolve_reference_title(chunk)
            references.append(
                ReferenceItem(
                    title=resolved_titles[file_hash],
                    url=f"graph://node/{file_hash}",
                    snippet=chunk["content"][:240],
                )
            )
        return references

    async def _normalize_chunks_for_prompt(
        self, chunks: List[RetrievedChunk]
    ) -> List[RetrievedChunk]:
        resolved_titles: Dict[str, str] = {}
        normalized_chunks: List[RetrievedChunk] = []
        for chunk in chunks:
            file_hash = chunk["file_hash"]
            if file_hash and file_hash not in resolved_titles:
                resolved_titles[file_hash] = await self._resolve_reference_title(chunk)
            normalized_chunk = RetrievedChunk(**chunk)
            if file_hash and file_hash in resolved_titles:
                normalized_chunk["title"] = resolved_titles[file_hash]
            normalized_chunks.append(normalized_chunk)
        return normalized_chunks

    @staticmethod
    def _context_digest(chunks: List[RetrievedChunk]) -> str:
        lines: List[str] = []
        for chunk in chunks:
            lines.append(
                (
                    f"Judul: {chunk['title']}\n"
                    f"Bagian: {chunk.get('metadata_header', '')}\n"
                    f"Isi: {chunk['content']}"
                )
            )
        return "\n\n".join(lines)

    @staticmethod
    def _sanitize_response(text: str) -> str:
        text = re.sub(
            r"(?is)\n*\*{0,2}\s*sumber\s+bacaan\s*:?\s*\*{0,2}.*$",
            "",
            text,
        )
        text = re.sub(r"\[(?:[a-f0-9]{32}|hash_dokumen|file_hash|document_hash)\]", "", text)
        text = re.sub(r"[ \t]+([,.;:!?])", r"\1", text)
        text = re.sub(r"([,.;:!?])(?!\s|$)", r"\1 ", text)
        text = re.sub(r"[ \t]+\n", "\n", text)
        text = re.sub(r"\n{3,}", "\n\n", text)
        return text.strip()

    @staticmethod
    def _append_markdown_links(response: str, references: List[ReferenceItem]) -> str:
        if not references:
            return response
        lines = ["", "", "**Sumber Bacaan:**"]
        for reference in references:
            lines.append(f"- [{reference['title']}]({reference['url']})")
        return response + "\n".join(lines)

    async def run(self, state: PRALState, writer: StreamWriter) -> PRALState:
        chunks = state["compressed_chunks"] or state["raw_chunks"]
        normalized_chunks = await self._normalize_chunks_for_prompt(chunks)
        references = await self._build_references(chunks)
        prompt = [
            {"role": "system", "content": self.policy.get_system_prompt()},
            {
                "role": "system",
                "content": self.policy.format_layer_instruction(state["layer"]),
            },
            {
                "role": "user",
                "content": (
                    f"Query: {state['query']}\n\n"
                    f"Konteks graf:\n{state['graph_context_text']}\n\n"
                    f"Konteks chunk:\n{self._context_digest(normalized_chunks)}\n\n"
                    "Gunakan markdown biasa. Referensi final akan dirender sebagai markdown links."
                ),
            },
        ]
        
        full_content = ""
        async for chunk in self._llm.astream(prompt):
            content = chunk.content
            if content:
                full_content += content
                writer({"token": content})

        cleaned_response = self._sanitize_response(full_content.strip())
        final_response = self._append_markdown_links(cleaned_response, references)
        state["references"] = references
        state["final_response"] = final_response
        state["metrics"]["reference_count"] = len(references)
        return state


class LangGraphPRALOrchestrator:
    """Runtime orchestrator that compiles and invokes the LangGraph PRAL flow."""

    def __init__(
        self,
        router_node: RouterNode | None = None,
        retrieval_flow: RetrievalFlow | None = None,
        compressor_node: CompressorNode | None = None,
        generator_node: GeneratorNode | None = None,
        history_manager: ChatHistoryManager | None = None,
    ) -> None:
        if StateGraph is None or START is None or END is None:
            raise RuntimeError("langgraph is not installed.")

        if retrieval_flow is None or generator_node is None:
            from app.services.chatbot.graph_retriever import GraphRetriever

            shared_graph_retriever = GraphRetriever()
            retrieval_flow = retrieval_flow or RetrievalFlow(
                graph_retriever=shared_graph_retriever
            )
            generator_node = generator_node or GeneratorNode(
                graph_retriever=shared_graph_retriever
            )

        self.router_node = router_node or RouterNode()
        self.retrieval_flow = retrieval_flow
        self.compressor_node = compressor_node or CompressorNode()
        self.generator_node = generator_node
        self.history_manager = history_manager or ChatHistoryManager()
        self.graph = self._build_graph()

    def _build_graph(self):
        workflow = StateGraph(PRALState)
        workflow.add_node("router", self.router_node.run)
        workflow.add_node("retrieval", self.retrieval_flow.run)
        workflow.add_node("compressor", self.compressor_node.run)
        workflow.add_node("generator", self.generator_node.run)
        workflow.add_edge(START, "router")
        workflow.add_conditional_edges(
            "router",
            self._route_after_router,
            {
                "retrieval": "retrieval",
                "generator": "generator",
            },
        )
        workflow.add_edge("retrieval", "compressor")
        workflow.add_edge("compressor", "generator")
        workflow.add_edge("generator", END)
        return workflow.compile()

    @staticmethod
    def _route_after_router(state: PRALState) -> str:
        return "generator" if state["layer"] == 3 else "retrieval"

    async def connect(self) -> None:
        await self.retrieval_flow.graph_retriever.connect()

    async def close(self) -> None:
        await self.retrieval_flow.graph_retriever.close()

    @staticmethod
    def _build_initial_state(
        query: str,
        active_hashes: List[str],
        project_id: str,
        user_id: str,
        mode: str,
    ) -> PRALState:
        normalized_mode: Literal["STRICT", "RESEARCH"] = (
            "RESEARCH" if mode == "RESEARCH" else "STRICT"
        )
        return PRALState(
            query=query,
            project_id=project_id,
            user_id=user_id,
            active_hashes=list(active_hashes),
            mode=normalized_mode,
            layer=1,
            intent="",
            target_hashes=list(active_hashes),
            neighbor_documents=[],
            graph_context_text="",
            raw_chunks=[],
            compressed_chunks=[],
            references=[],
            final_response="",
            metrics={
                "project_id": project_id,
                "layer_strategy": 1,
                "query_intent": "",
                "source_count": 0,
            },
        )

    @staticmethod
    def _strict_empty_message(has_used_research: bool) -> str:
        if has_used_research:
            return (
                "Saat ini Anda berada di mode **STRICT**. Mode ini membutuhkan setidaknya "
                "satu dokumen aktif untuk dianalisis. Silakan pilih dokumen melalui "
                "*Canvas Knowledge Graph* terlebih dahulu.\n\nAtau, jika Anda ingin "
                "mencari hubungan antar-literatur secara bebas, silakan alihkan kembali "
                "ke mode **RESEARCH**."
            )
        return (
            "Untuk memulai analisis, silakan pilih dokumen melalui *Canvas Knowledge Graph* "
            "atau gunakan ikon tambah (+). Jika Anda belum memiliki literatur dalam sistem, "
            "silakan unggah dokumen PDF Anda terlebih dahulu. Setelah dokumen aktif, Anda "
            "dapat mengajukan pertanyaan terkait konten literatur tersebut."
        )

    async def process_chat(
        self,
        query: str,
        active_hashes: List[str],
        project_id: str,
        user_id: str,
        mode: str = "STRICT",
        has_used_research: bool = False,
    ) -> AsyncGenerator[str, None]:
        start_time = time.time()

        if mode == "STRICT" and not active_hashes:
            message = self._strict_empty_message(has_used_research)
            yield message
            return

        state = self._build_initial_state(
            query=query,
            active_hashes=active_hashes,
            project_id=project_id,
            user_id=user_id,
            mode=mode,
        )

        try:
            final_state = state
            async for stream_mode, payload in self.graph.astream(state, stream_mode=["updates", "custom"]):
                if stream_mode == "custom":
                    token = payload.get("token", "")
                    if token:
                        yield json.dumps({
                            "event": "chunk",
                            "data": json.dumps({"text": token})
                        }) + "\n\n"
                elif stream_mode == "updates":
                    for node_name, node_state in payload.items():
                        final_state = node_state
                        
                        if node_name == "router":
                            intent = node_state.get("intent", "Routing")
                            yield json.dumps({
                                "event": "milestone",
                                "data": json.dumps({"step": "router", "label": intent})
                            }) + "\n\n"
                        elif node_name == "retrieval":
                            raw_chunks = node_state.get("raw_chunks", [])
                            yield json.dumps({
                                "event": "milestone",
                                "data": json.dumps({"step": "retrieval", "label": f"Found {len(raw_chunks)} chunks"})
                            }) + "\n\n"
                        elif node_name == "compressor":
                            compressed = node_state.get("compressed_chunks", [])
                            yield json.dumps({
                                "event": "milestone",
                                "data": json.dumps({"step": "compressor", "label": f"Compressed to {len(compressed)} chunks"})
                            }) + "\n\n"
                        elif node_name == "generator":
                            pass

            latency_ms = int((time.time() - start_time) * 1000)
            final_state["metrics"]["latency_ms"] = latency_ms

            self.history_manager.save_turn(
                project_id=project_id,
                user_id=user_id,
                user_query=query,
                ai_response=final_state.get("final_response", ""),
                references=final_state.get("references", []),
                metadata=final_state.get("metrics", {}),
            )
            
            yield json.dumps({
                "event": "final",
                "data": json.dumps({
                    "references": final_state.get("references", []),
                    "metrics": final_state.get("metrics", {})
                })
            }) + "\n\n"
        except Exception as exc:
            logger.error("LangGraph orchestrator failed: %s", exc)
            yield json.dumps({
                "event": "error",
                "data": json.dumps({
                    "message": f"Maaf, terjadi kesalahan saat memproses respons: {str(exc)}",
                    "code": 500
                })
            }) + "\n\n"
