"""
LangGraph core span-tracking engine.

Tracks active runs (graph, node, LLM, tool, retriever) by run ID,
records timing and metadata, and produces span dicts when runs complete.
Used by LangGraphHandler to decouple callback protocol from span construction.
"""

import logging
import uuid
from datetime import datetime, timezone
from typing import Any, Dict, List

from .cost_tracking import get_langgraph_cost
from .error_detection import get_error_detector

logger = logging.getLogger(__name__)


class _ActiveRun:
    """Tracks a single active run (node, LLM call, tool call, etc.)."""

    __slots__ = (
        "input",
        "metadata",
        "name",
        "parent_run_id",
        "run_id",
        "run_type",
        "start_time",
        "tokens",
    )

    def __init__(
        self,
        run_id: str,
        parent_run_id: str | None,
        run_type: str,
        name: str,
        input: Any = None,
        metadata: Dict[str, Any] | None = None,
    ):
        self.run_id = run_id
        self.parent_run_id = parent_run_id
        self.run_type = run_type
        self.name = name
        self.input = input
        self.start_time = datetime.now(timezone.utc)
        self.metadata = metadata or {}
        self.tokens: str = ""  # For streaming LLM token accumulation


class LangGraphCore:
    """Stateful span tracker for LangGraph executions.

    Maintains a registry of active runs keyed by run_id.
    Start methods register a run; end/error methods pop it and return a span dict.
    """

    def __init__(self):
        self._active_runs: Dict[str, _ActiveRun] = {}
        self._trace_runs: Dict[str, _ActiveRun] = {}
        self._error_detector = get_error_detector()

    # ------------------------------------------------------------------
    # Trace lifecycle
    # ------------------------------------------------------------------

    def start_trace(
        self,
        run_id: str,
        trace_obj: Any,
        trace_name: str,
        input: Any,
    ) -> None:
        """Register the root graph execution."""
        run = _ActiveRun(
            run_id=run_id,
            parent_run_id=None,
            run_type="graph",
            name=trace_name,
            input=input,
            metadata={
                "trace_id": getattr(trace_obj, "id", str(trace_obj)),
                "framework": "langgraph",
            },
        )
        self._active_runs[run_id] = run
        self._trace_runs[run_id] = run

    def end_trace(self, run_id: str, output: Any) -> List[Dict[str, Any]]:
        """Finalize the root graph execution. Returns list of event dicts."""
        run = self._trace_runs.pop(run_id, None) or self._active_runs.pop(run_id, None)
        if not run:
            return []

        now = datetime.now(timezone.utc)
        duration_ms = (now - run.start_time).total_seconds() * 1000
        trace_id = run.metadata.get("trace_id", run_id)

        workflow_span = {
            "id": str(uuid.uuid4()),
            "trace_id": trace_id,
            "name": run.name,
            "type": "chain",
            "input": _safe_serialize(run.input),
            "output": _safe_serialize(output),
            "status": "completed",
            "start_time": run.start_time.isoformat(),
            "end_time": now.isoformat(),
            "metadata": {
                **run.metadata,
                "duration_ms": round(duration_ms, 2),
                "nodeType": "workflow",
            },
        }

        trace_update = {
            "id": trace_id,
            "status": "completed",
            "end_time": now.isoformat(),
            "metadata": {"duration_ms": round(duration_ms, 2)},
        }

        return [
            {"type": "span-create", "body": workflow_span},
            {"type": "trace-update", "body": trace_update},
        ]

    def error_trace(self, run_id: str, error: Exception) -> List[Dict[str, Any]]:
        """Finalize a failed graph execution."""
        run = self._trace_runs.pop(run_id, None) or self._active_runs.pop(run_id, None)
        if not run:
            return []

        now = datetime.now(timezone.utc)
        duration_ms = (now - run.start_time).total_seconds() * 1000
        trace_id = run.metadata.get("trace_id", run_id)

        detected = self._error_detector.detect_from_exception(error, f"graph:{run.name}", {})

        workflow_span = {
            "id": str(uuid.uuid4()),
            "trace_id": trace_id,
            "name": run.name,
            "type": "chain",
            "input": _safe_serialize(run.input),
            "output": None,
            "status": "error",
            "error": str(error),
            "error_type": type(error).__name__,
            "start_time": run.start_time.isoformat(),
            "end_time": now.isoformat(),
            "metadata": {
                **run.metadata,
                "duration_ms": round(duration_ms, 2),
                "nodeType": "workflow",
                **(
                    {
                        "error_detection": {
                            "type": detected.error_type.value,
                            "severity": detected.severity.value,
                            "is_transient": detected.is_transient,
                        }
                    }
                    if detected
                    else {}
                ),
            },
        }

        trace_update = {
            "id": trace_id,
            "status": "failed",
            "error": str(error),
            "end_time": now.isoformat(),
        }

        return [
            {"type": "span-create", "body": workflow_span},
            {"type": "trace-update", "body": trace_update},
        ]

    # ------------------------------------------------------------------
    # Node lifecycle
    # ------------------------------------------------------------------

    def start_node(
        self,
        run_id: str,
        parent_run_id: str | None,
        node_name: str,
        input: Any,
    ) -> None:
        """Register a node execution start."""
        self._active_runs[run_id] = _ActiveRun(
            run_id=run_id,
            parent_run_id=parent_run_id,
            run_type="node",
            name=node_name,
            input=input,
        )

    def end_node(self, run_id: str, output: Any) -> Dict[str, Any] | None:
        """Finalize a node execution. Returns span dict or None."""
        run = self._active_runs.pop(run_id, None)
        if not run:
            return None
        return self._build_span(run, output=output, status="completed")

    def error_node(self, run_id: str, error: Exception) -> Dict[str, Any] | None:
        """Finalize a failed node execution."""
        run = self._active_runs.pop(run_id, None)
        if not run:
            return None

        detected = self._error_detector.detect_from_exception(error, f"node:{run.name}", {})
        return self._build_span(run, output=None, status="error", error=error, detected=detected)

    # ------------------------------------------------------------------
    # LLM lifecycle
    # ------------------------------------------------------------------

    def start_llm(
        self,
        run_id: str,
        parent_run_id: str | None,
        model_name: str,
        prompts: Any,
    ) -> None:
        """Register an LLM call start."""
        self._active_runs[run_id] = _ActiveRun(
            run_id=run_id,
            parent_run_id=parent_run_id,
            run_type="llm",
            name=model_name,
            input=prompts,
            metadata={"model": model_name},
        )

    def end_llm(self, run_id: str, response: Any) -> Dict[str, Any] | None:
        """Finalize an LLM call. Returns span dict with token/cost data."""
        run = self._active_runs.pop(run_id, None)
        if not run:
            return None

        # Extract token usage from LangChain LLMResult
        usage = {}
        try:
            if hasattr(response, "llm_output") and response.llm_output:
                token_usage = response.llm_output.get("token_usage", {})
                usage = {
                    "input_tokens": token_usage.get("prompt_tokens", 0),
                    "output_tokens": token_usage.get("completion_tokens", 0),
                    "total_tokens": token_usage.get("total_tokens", 0),
                }
            elif hasattr(response, "generations") and response.generations:
                for gen_list in response.generations:
                    for gen in gen_list:
                        gen_info = getattr(gen, "generation_info", {}) or {}
                        if "usage_metadata" in gen_info:
                            um = gen_info["usage_metadata"]
                            usage = {
                                "input_tokens": um.get("input_tokens", 0)
                                or um.get("prompt_tokens", 0),
                                "output_tokens": um.get("output_tokens", 0)
                                or um.get("completion_tokens", 0),
                                "total_tokens": um.get("total_tokens", 0),
                            }
                            break
                    if usage:
                        break
        except Exception:
            pass

        # Estimate cost
        cost = None
        model = run.metadata.get("model", "")
        if usage and model:
            try:
                cost_data = get_langgraph_cost(
                    model,
                    usage.get("input_tokens", 0),
                    usage.get("output_tokens", 0),
                )
                if cost_data:
                    cost = cost_data.get("total_cost", 0)
            except Exception:
                pass

        span = self._build_span(run, output=_safe_serialize(response), status="completed")
        if span:
            span["type"] = "llm"
            if usage:
                span["metadata"]["usage"] = usage
            if cost is not None:
                span["metadata"]["cost"] = {"total_cost": cost}
        return span

    def llm_new_token(self, run_id: str, token: str) -> None:
        """Accumulate a streaming token."""
        run = self._active_runs.get(run_id)
        if run:
            run.tokens += token

    # ------------------------------------------------------------------
    # Tool lifecycle
    # ------------------------------------------------------------------

    def start_tool(
        self,
        run_id: str,
        parent_run_id: str | None,
        tool_name: str,
        input_str: Any,
    ) -> None:
        """Register a tool call start."""
        self._active_runs[run_id] = _ActiveRun(
            run_id=run_id,
            parent_run_id=parent_run_id,
            run_type="tool",
            name=tool_name,
            input=input_str,
        )

    def end_tool(self, run_id: str, output: Any) -> Dict[str, Any] | None:
        """Finalize a tool call."""
        run = self._active_runs.pop(run_id, None)
        if not run:
            return None

        is_error = _is_error_output(output)
        status = "error" if is_error else "completed"
        span = self._build_span(run, output=output, status=status)
        if span and is_error:
            span["is_error"] = True
            span["error"] = str(output)
        return span

    # ------------------------------------------------------------------
    # Retriever lifecycle
    # ------------------------------------------------------------------

    def start_retriever(
        self,
        run_id: str,
        parent_run_id: str | None,
        query: str,
        serialized: Any,
    ) -> None:
        """Register a retriever call start."""
        name = "Retriever"
        if serialized and isinstance(serialized, dict):
            name = serialized.get("name", "Retriever")
        self._active_runs[run_id] = _ActiveRun(
            run_id=run_id,
            parent_run_id=parent_run_id,
            run_type="retriever",
            name=name,
            input=query,
        )

    def end_retriever(self, run_id: str, documents: Any) -> Dict[str, Any] | None:
        """Finalize a retriever call."""
        run = self._active_runs.pop(run_id, None)
        if not run:
            return None

        doc_count = len(documents) if documents else 0
        span = self._build_span(run, output=f"{doc_count} documents", status="completed")
        if span:
            span["metadata"]["document_count"] = doc_count
        return span

    def error_retriever(self, run_id: str, error: BaseException) -> Dict[str, Any] | None:
        """Finalize a failed retriever call."""
        run = self._active_runs.pop(run_id, None)
        if not run:
            return None

        detected = self._error_detector.detect_from_exception(
            error if isinstance(error, Exception) else Exception(str(error)),
            f"retriever:{run.name}",
            {},
        )
        return self._build_span(run, output=None, status="error", error=error, detected=detected)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _build_span(
        self,
        run: _ActiveRun,
        *,
        output: Any,
        status: str,
        error: BaseException | None = None,
        detected: Any = None,
    ) -> Dict[str, Any]:
        """Build a span dict from a completed run."""
        now = datetime.now(timezone.utc)
        duration_ms = (now - run.start_time).total_seconds() * 1000

        # Find trace_id from parent chain
        trace_id = run.metadata.get("trace_id")
        if not trace_id and run.parent_run_id:
            parent = self._active_runs.get(run.parent_run_id) or self._trace_runs.get(
                run.parent_run_id
            )
            if parent:
                trace_id = parent.metadata.get("trace_id")

        span = {
            "id": str(uuid.uuid4()),
            "trace_id": trace_id or "",
            "parent_span_id": run.parent_run_id,
            "name": run.name,
            "type": run.run_type,
            "input": _safe_serialize(run.input),
            "output": _safe_serialize(output),
            "status": status,
            "start_time": run.start_time.isoformat(),
            "end_time": now.isoformat(),
            "metadata": {
                **run.metadata,
                "duration_ms": round(duration_ms, 2),
                "nodeType": run.run_type,
            },
        }

        if error:
            span["error"] = str(error)
            span["error_type"] = type(error).__name__

        if detected:
            span["metadata"]["error_detection"] = {
                "type": detected.error_type.value,
                "severity": detected.severity.value,
                "is_transient": detected.is_transient,
            }

        return span


# ------------------------------------------------------------------
# Module-level helpers
# ------------------------------------------------------------------


def _safe_serialize(value: Any) -> Any:
    """Safely convert a value to something JSON-serializable."""
    if value is None:
        return None
    if isinstance(value, (str, int, float, bool)):
        return value
    if isinstance(value, dict):
        return {k: _safe_serialize(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [_safe_serialize(v) for v in value]
    try:
        return str(value)[:2000]
    except Exception:
        return "<unserializable>"


def _is_error_output(output: Any) -> bool:
    """Check if a tool output represents an error."""
    if output is None:
        return False
    output_str = str(output).lower()
    error_indicators = ("error:", "exception:", "traceback", "failed:", "status: error")
    return any(indicator in output_str for indicator in error_indicators)
