"""
Native LangSmith tracer for the extraction agents and the playground's
RAG agents. Each pydantic-ai agent class in `quber.agents` holds its own
instance. The playground's agents share one process-wide instance from
`quber.playground.tracing`. The cli backend (`ClaudeCLIClient`) and the
`quber.processors` classes built through `AgentFactory` are not traced.

Posts runs to LangSmith's first-class `/runs/multipart` API with explicit
`inputs` and `outputs` JSON — the same pattern as
`~/.claude/hooks/stop_hook.sh`. We do this instead of OTLP because the
OTel adapter on LangSmith's side has to introspect span attributes to
derive inputs/outputs, and Logfire/pydantic-ai emit attribute keys that
the adapter doesn't recognize. The result: traces land, but the LangSmith
UI shows empty input/output panels.

This tracer skips the translation entirely. We hand LangSmith the data
in the shape its UI renders directly.

Toggle via `TRACE_TO_LANGSMITH` env var; configured via
`CC_LANGSMITH_API_KEY`, `CC_LANGSMITH_PROJECT`, `CC_LANGSMITH_DEBUG`.
"""

from __future__ import annotations

import re
import uuid
from contextlib import asynccontextmanager, contextmanager
from datetime import datetime, timezone
from typing import Any, Dict, Generator, Optional

from loguru import logger

from quber.settings import get_settings


def langsmith_enabled() -> bool:
    obs = get_settings().obs
    return obs.trace_to_langsmith and bool(obs.langsmith_api_key)


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


_DATED_SUFFIX = re.compile(r"-\d{8}$")


def canonical_model_name(model: str) -> str:
    """Strip the trailing `-YYYYMMDD` dated suffix from an Anthropic
    model identifier.

    LangSmith's pricing catalog and many of its UI affordances key on
    the bare model name (e.g. `claude-haiku-4-5`, `claude-sonnet-4-6`).
    The dated form Anthropic returns (`claude-haiku-4-5-20251001`)
    doesn't match the catalog, so cost attribution silently breaks.

    Same approach as `~/.claude/hooks/stop_hook.sh`'s sed.
    """
    return _DATED_SUFFIX.sub("", model or "")


def usage_metadata_from(usage: Any) -> Dict[str, Any]:
    """Translate a pydantic-ai `RunUsage` into LangSmith's `usage_metadata`
    schema (the LangChain standard).

    The reported `input_tokens` is the usage object's `input_tokens` plus
    its `cache_read_tokens` and `cache_write_tokens`. A pydantic-ai
    `RunUsage` built from an Anthropic response already counts cache reads
    and writes in `input_tokens`, so with prompt caching active the cached
    input is counted twice.
    """
    input_tokens = (
        (getattr(usage, "input_tokens", 0) or 0)
        + (getattr(usage, "cache_read_tokens", 0) or 0)
        + (getattr(usage, "cache_write_tokens", 0) or 0)
    )
    return {
        "input_tokens": int(input_tokens),
        "output_tokens": int(getattr(usage, "output_tokens", 0) or 0),
        "input_token_details": {
            "cache_read": int(getattr(usage, "cache_read_tokens", 0) or 0),
            "cache_creation": int(getattr(usage, "cache_write_tokens", 0) or 0),
        },
    }


class LangSmithTracer:
    """Authors LangSmith runs directly via the SDK. Each pydantic-ai agent
    class in `quber.agents` holds its own instance. The playground's agents
    share one from `quber.playground.tracing`.

    The tracer is a no-op if `TRACE_TO_LANGSMITH` is unset/false or
    `CC_LANGSMITH_API_KEY` is missing — callers don't need to guard.
    """

    def __init__(
        self,
        project: Optional[str] = None,
        api_key: Optional[str] = None,
        run_name: Optional[str] = None,
    ) -> None:
        self.enabled = langsmith_enabled()
        if not self.enabled:
            self.client = None
            self.project = None
            self.run_name = None
            return

        from langsmith import Client

        obs = get_settings().obs
        self.client = Client(
            api_url="https://api.smith.langchain.com",
            api_key=api_key or obs.langsmith_api_key,
        )
        self.project = project or obs.langsmith_project
        self.run_name = run_name or "quber"

        if obs.langsmith_debug:
            logger.info(
                "LangSmith tracer initialized: project={} run_name={}",
                self.project,
                self.run_name,
            )

    def build_extra(
        self,
        model: Optional[str],
        extra_metadata: Optional[Dict[str, Any]],
    ) -> Dict[str, Any]:
        canonical = canonical_model_name(model) if model else None
        metadata: Dict[str, Any] = {"ls_provider": "anthropic"}
        if canonical:
            metadata["ls_model_name"] = canonical
        if model and model != canonical:
            metadata["ls_model_full_name"] = model
        if extra_metadata:
            metadata.update(extra_metadata)
        return {"metadata": metadata}

    @contextmanager
    def llm_run_sync(
        self,
        name: str,
        inputs: Dict[str, Any],
        model: Optional[str] = None,
        extra_metadata: Optional[Dict[str, Any]] = None,
    ) -> Generator["RunHandle"]:
        """Open an LLM-type run. Caller sets `.outputs` on the yielded
        handle, and the tracer patches LangSmith on exit with the outputs
        and the end timestamp. The run is recorded as failed only when an
        exception escapes the `with` block. The tracer records its `repr`
        and re-raises it.
        """
        handle = RunHandle()
        if not self.enabled or self.client is None:
            yield handle
            return

        run_id = str(uuid.uuid4())
        start_time = utc_now()
        canonical = canonical_model_name(model) if model else None
        extra = self.build_extra(model, extra_metadata)

        try:
            self.client.create_run(
                id=run_id,
                name=name,
                run_type="llm",
                inputs=inputs,
                start_time=start_time,
                project_name=self.project,
                extra=extra,
                tags=[canonical] if canonical else None,
            )
        except Exception as exc:
            logger.warning("LangSmith create_run failed: {}", exc)
            yield handle
            return

        try:
            yield handle
        except Exception as exc:
            try:
                self.client.update_run(
                    run_id=run_id,
                    error=repr(exc),
                    end_time=utc_now(),
                )
            except Exception as patch_exc:
                logger.warning("LangSmith update_run (error) failed: {}", patch_exc)
            raise
        else:
            try:
                self.client.update_run(
                    run_id=run_id,
                    outputs=handle.outputs or {},
                    end_time=utc_now(),
                )
            except Exception as exc:
                logger.warning("LangSmith update_run (ok) failed: {}", exc)

    @asynccontextmanager
    async def llm_run(
        self,
        name: str,
        inputs: Dict[str, Any],
        model: Optional[str] = None,
        extra_metadata: Optional[Dict[str, Any]] = None,
    ):
        """Async sibling of `llm_run_sync`.

        The langsmith Client's `create_run` / `update_run` are non-blocking
        in practice (they enqueue onto a background batched sender), so we
        can call them from async code without awaiting. The trace-side
        I/O happens off the request path.
        """
        handle = RunHandle()
        if not self.enabled or self.client is None:
            yield handle
            return

        run_id = str(uuid.uuid4())
        start_time = utc_now()
        canonical = canonical_model_name(model) if model else None
        extra = self.build_extra(model, extra_metadata)

        try:
            self.client.create_run(
                id=run_id,
                name=name,
                run_type="llm",
                inputs=inputs,
                start_time=start_time,
                project_name=self.project,
                extra=extra,
                tags=[canonical] if canonical else None,
            )
        except Exception as exc:
            logger.warning("LangSmith create_run failed: {}", exc)
            yield handle
            return

        try:
            yield handle
        except Exception as exc:
            try:
                self.client.update_run(run_id=run_id, error=repr(exc), end_time=utc_now())
            except Exception as patch_exc:
                logger.warning("LangSmith update_run (error) failed: {}", patch_exc)
            raise
        else:
            try:
                self.client.update_run(run_id=run_id, outputs=handle.outputs or {}, end_time=utc_now())
            except Exception as exc:
                logger.warning("LangSmith update_run (ok) failed: {}", exc)


class RunHandle:
    """Mutable handle the caller sets `outputs` on inside the context."""

    __slots__ = ("outputs",)

    def __init__(self) -> None:
        self.outputs: Dict[str, Any] = {}