from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
import json
from typing import Any
from uuid import UUID, uuid4

from agenttrace.trace import trace


def _now_utc() -> datetime:
    return datetime.now(timezone.utc)


def _to_run_key(value: UUID | str | None) -> str | None:
    if value is None:
        return None
    return str(value)


def _pick_name(serialized: dict[str, Any] | None, fallback: str) -> str:
    if not serialized:
        return fallback
    value = serialized.get("name")
    if isinstance(value, str) and value:
        return value
    return fallback


def _safe_json(value: Any, seen: set[int] | None = None) -> Any:
    if seen is None:
        seen = set()
    if value is None or isinstance(value, (str, int, float, bool)):
        return value

    value_id = id(value)
    if value_id in seen:
        return "<recursive>"

    if isinstance(value, dict):
        seen.add(value_id)
        return {str(k): _safe_json(v, seen) for k, v in value.items()}
    if isinstance(value, (list, tuple, set)):
        seen.add(value_id)
        return [_safe_json(v, seen) for v in value]
    if isinstance(value, bytes):
        return value.decode("utf-8", errors="replace")

    for attr in ("model_dump", "dict"):
        fn = getattr(value, attr, None)
        if callable(fn):
            try:
                return _safe_json(fn(), seen)
            except Exception:  # noqa: BLE001
                pass

    try:
        json.dumps(value)
        return value
    except TypeError:
        return repr(value)


@dataclass(slots=True)
class _PendingSpan:
    span_id: str
    span_type: str
    name: str
    start_time: datetime
    input_json: Any
    metadata_json: dict[str, Any]
    parent_run_key: str | None


class AgentTraceCallbackHandler:
    def __init__(self) -> None:
        # LangChain callback manager expects these handler attributes.
        self.raise_error = False
        self.run_inline = True
        self.ignore_llm = False
        self.ignore_chat_model = False
        self.ignore_chain = False
        self.ignore_agent = False
        self.ignore_retriever = False
        self.ignore_tool = False
        self.ignore_custom_event = False
        self.ignore_retry = False

        self._pending: dict[str, _PendingSpan] = {}
        self._span_id_by_run: dict[str, str] = {}

    def _start(
        self,
        *,
        span_type: str,
        name: str,
        input_json: Any,
        run_id: UUID,
        parent_run_id: UUID | None,
        metadata_json: dict[str, Any] | None = None,
    ) -> None:
        run_key = _to_run_key(run_id)
        span_id = str(uuid4())
        self._pending[run_key] = _PendingSpan(
            span_id=span_id,
            span_type=span_type,
            name=name,
            start_time=_now_utc(),
            input_json=_safe_json(input_json),
            metadata_json=_safe_json(metadata_json or {}),
            parent_run_key=_to_run_key(parent_run_id),
        )
        self._span_id_by_run[run_key] = span_id

    def _end(self, *, run_id: UUID, output_json: Any, error: str | None = None) -> None:
        ctx = trace.current()
        if ctx is None:
            return
        run_key = _to_run_key(run_id)
        pending = self._pending.pop(run_key, None)
        end_time = _now_utc()
        if pending is None:
            pending = _PendingSpan(
                span_id=str(uuid4()),
                span_type="chain",
                name="unknown",
                start_time=end_time,
                input_json={},
                metadata_json={},
                parent_run_key=None,
            )

        parent_span_id = None
        if pending.parent_run_key:
            parent_span_id = self._span_id_by_run.get(pending.parent_run_key)

        span = ctx.add_span(
            span_id=pending.span_id,
            span_type=pending.span_type,  # type: ignore[arg-type]
            name=pending.name,
            start_time=pending.start_time,
            end_time=end_time,
            input_json=pending.input_json,
            output_json=_safe_json(output_json),
            error=error,
            metadata_json=pending.metadata_json,
            parent_span_id=parent_span_id,
        )
        self._span_id_by_run[run_key] = str(span.span_id)

    def on_llm_start(
        self,
        serialized: dict[str, Any],
        prompts: list[str],
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **_: Any,
    ) -> None:
        self._start(
            span_type="llm",
            name=_pick_name(serialized, "llm"),
            input_json={"prompts": prompts},
            run_id=run_id,
            parent_run_id=parent_run_id,
            metadata_json={"serialized": serialized},
        )

    def on_llm_end(self, response: Any, *, run_id: UUID, **_: Any) -> None:
        self._end(run_id=run_id, output_json={"response": response})

    def on_chat_model_start(
        self,
        serialized: dict[str, Any],
        messages: list[list[Any]],
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **_: Any,
    ) -> None:
        self._start(
            span_type="llm",
            name=_pick_name(serialized, "chat_model"),
            input_json={"messages": messages},
            run_id=run_id,
            parent_run_id=parent_run_id,
            metadata_json={"serialized": serialized},
        )

    def on_chat_model_end(self, response: Any, *, run_id: UUID, **_: Any) -> None:
        self._end(run_id=run_id, output_json={"response": response})

    def on_llm_new_token(self, token: str, *, run_id: UUID, **_: Any) -> None:
        # Token-level streaming is not persisted in v0.
        return

    def on_llm_error(self, error: BaseException, *, run_id: UUID, **_: Any) -> None:
        self._end(
            run_id=run_id,
            output_json={},
            error=f"{type(error).__name__}: {error}",
        )

    def on_tool_start(
        self,
        serialized: dict[str, Any],
        input_str: str,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **_: Any,
    ) -> None:
        self._start(
            span_type="tool",
            name=_pick_name(serialized, "tool"),
            input_json={"input": input_str},
            run_id=run_id,
            parent_run_id=parent_run_id,
            metadata_json={"serialized": serialized},
        )

    def on_tool_end(self, output: Any, *, run_id: UUID, **_: Any) -> None:
        self._end(run_id=run_id, output_json={"output": output})

    def on_tool_error(self, error: BaseException, *, run_id: UUID, **_: Any) -> None:
        self._end(
            run_id=run_id,
            output_json={},
            error=f"{type(error).__name__}: {error}",
        )

    def on_retriever_start(
        self,
        serialized: dict[str, Any],
        query: str,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **_: Any,
    ) -> None:
        self._start(
            span_type="retrieval",
            name=_pick_name(serialized, "retriever"),
            input_json={"query": query},
            run_id=run_id,
            parent_run_id=parent_run_id,
            metadata_json={"serialized": serialized},
        )

    def on_retriever_end(self, documents: Any, *, run_id: UUID, **_: Any) -> None:
        self._end(run_id=run_id, output_json={"documents": documents})

    def on_retriever_error(self, error: BaseException, *, run_id: UUID, **_: Any) -> None:
        self._end(
            run_id=run_id,
            output_json={},
            error=f"{type(error).__name__}: {error}",
        )

    def on_chain_start(
        self,
        serialized: dict[str, Any],
        inputs: Any,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **_: Any,
    ) -> None:
        self._start(
            span_type="chain",
            name=_pick_name(serialized, "chain"),
            input_json=inputs,
            run_id=run_id,
            parent_run_id=parent_run_id,
            metadata_json={"serialized": serialized},
        )

    def on_chain_end(self, outputs: dict[str, Any], *, run_id: UUID, **_: Any) -> None:
        self._end(run_id=run_id, output_json=outputs)

    def on_chain_error(self, error: BaseException, *, run_id: UUID, **_: Any) -> None:
        self._end(
            run_id=run_id,
            output_json={},
            error=f"{type(error).__name__}: {error}",
        )
