"""LangGraph adapter — instrument your graph in 2 lines.

Usage:
    from velenai_sdk import VelenAI
    from velenai_sdk.adapters.langgraph import VelenAILangGraphCallback

    vai = VelenAI()

    # Option 1: As a LangChain callback
    graph = StateGraph(...)
    result = graph.compile().invoke(input, config={"callbacks": [VelenAILangGraphCallback(vai)]})

    # Option 2: Wrap the graph
    from velenai_sdk.adapters.langgraph import instrument_graph
    instrumented = instrument_graph(vai, graph.compile(), agent_id="my-graph")
    result = instrumented.invoke(input)
"""

from __future__ import annotations

import time
from typing import Any

from ..client import VelenAI


class VelenAILangGraphCallback:
    """LangChain/LangGraph callback handler.

    Tracks each LLM call and chain run, emitting telemetry to VelenAI.
    Compatible with LangChain's BaseCallbackHandler interface.
    """

    def __init__(self, client: VelenAI, agent_prefix: str = "langgraph"):
        self._client = client
        self._prefix = agent_prefix
        self._run_starts: dict[str, float] = {}
        self._run_tokens: dict[str, int] = {}

    # --- LLM callbacks ---

    def on_llm_start(self, serialized: dict, prompts: list, *, run_id: Any, **kwargs: Any) -> None:
        self._run_starts[str(run_id)] = time.monotonic()

    def on_llm_end(self, response: Any, *, run_id: Any, **kwargs: Any) -> None:
        rid = str(run_id)
        start = self._run_starts.pop(rid, time.monotonic())
        latency_ms = int((time.monotonic() - start) * 1000)

        tokens = None
        if hasattr(response, "llm_output") and response.llm_output:
            usage = response.llm_output.get("token_usage", {})
            tokens = usage.get("total_tokens")

        self._client.emit(
            agent_id=f"{self._prefix}.llm",
            success=True,
            latency_ms=latency_ms,
            tokens=tokens,
            metadata={"framework": "langgraph", "type": "llm_call"},
        )

    def on_llm_error(self, error: Exception, *, run_id: Any, **kwargs: Any) -> None:
        rid = str(run_id)
        start = self._run_starts.pop(rid, time.monotonic())
        latency_ms = int((time.monotonic() - start) * 1000)

        self._client.emit(
            agent_id=f"{self._prefix}.llm",
            success=False,
            latency_ms=latency_ms,
            metadata={"framework": "langgraph", "error": str(error)[:500]},
        )

    # --- Chain/graph callbacks ---

    def on_chain_start(self, serialized: dict, inputs: dict, *, run_id: Any, **kwargs: Any) -> None:
        self._run_starts[str(run_id)] = time.monotonic()

    def on_chain_end(self, outputs: dict, *, run_id: Any, **kwargs: Any) -> None:
        rid = str(run_id)
        start = self._run_starts.pop(rid, time.monotonic())
        latency_ms = int((time.monotonic() - start) * 1000)

        name = "chain"
        if "serialized" in (outputs or {}):
            name = outputs.get("name", "chain")

        self._client.emit(
            agent_id=f"{self._prefix}.{name}",
            success=True,
            latency_ms=latency_ms,
            metadata={"framework": "langgraph", "type": "chain_run"},
        )

    def on_chain_error(self, error: Exception, *, run_id: Any, **kwargs: Any) -> None:
        rid = str(run_id)
        start = self._run_starts.pop(rid, time.monotonic())
        latency_ms = int((time.monotonic() - start) * 1000)

        self._client.emit(
            agent_id=f"{self._prefix}.chain",
            success=False,
            latency_ms=latency_ms,
            metadata={"framework": "langgraph", "error": str(error)[:500]},
        )

    # --- Tool callbacks ---

    def on_tool_start(self, serialized: dict, input_str: str, *, run_id: Any, **kwargs: Any) -> None:
        self._run_starts[str(run_id)] = time.monotonic()

    def on_tool_end(self, output: str, *, run_id: Any, **kwargs: Any) -> None:
        rid = str(run_id)
        start = self._run_starts.pop(rid, time.monotonic())
        latency_ms = int((time.monotonic() - start) * 1000)

        self._client.emit(
            agent_id=f"{self._prefix}.tool",
            success=True,
            latency_ms=latency_ms,
            metadata={"framework": "langgraph", "type": "tool_call"},
        )

    def on_tool_error(self, error: Exception, *, run_id: Any, **kwargs: Any) -> None:
        rid = str(run_id)
        start = self._run_starts.pop(rid, time.monotonic())
        latency_ms = int((time.monotonic() - start) * 1000)

        self._client.emit(
            agent_id=f"{self._prefix}.tool",
            success=False,
            latency_ms=latency_ms,
            metadata={"framework": "langgraph", "error": str(error)[:500]},
        )


def instrument_graph(client: VelenAI, compiled_graph: Any, agent_id: str = "langgraph") -> Any:
    """Wrap a compiled LangGraph so every invoke/ainvoke is tracked.

    Usage:
        graph = StateGraph(...).compile()
        graph = instrument_graph(vai, graph, agent_id="my-workflow")
        result = graph.invoke(input)
    """

    class InstrumentedGraph:
        def __init__(self, inner, vai, aid):
            self._inner = inner
            self._vai = vai
            self._aid = aid
            # Proxy all attributes
            for attr in dir(inner):
                if not attr.startswith("_") and attr not in ("invoke", "ainvoke"):
                    try:
                        setattr(self, attr, getattr(inner, attr))
                    except Exception:
                        pass

        def invoke(self, *args, **kwargs):
            with self._vai.trace(self._aid, framework="langgraph") as t:
                try:
                    result = self._inner.invoke(*args, **kwargs)
                    t.success = True
                    return result
                except Exception:
                    t.success = False
                    raise

        async def ainvoke(self, *args, **kwargs):
            start = time.monotonic()
            success = True
            try:
                result = await self._inner.ainvoke(*args, **kwargs)
                return result
            except Exception:
                success = False
                raise
            finally:
                await self._vai.emit_async(
                    agent_id=self._aid,
                    success=success,
                    latency_ms=int((time.monotonic() - start) * 1000),
                    metadata={"framework": "langgraph"},
                )

    return InstrumentedGraph(compiled_graph, client, agent_id)
