"""High-level conversation tracing context managers.

Provides the Turn-based API for recording conversation messages:

- ``TurnContext`` — a single conversational turn (user or assistant)
- ``ConversationContext`` — wraps a conversation span
- ``ToolsContext`` / ``ToolSpan`` — tool batch execution spans
- ``conversation()`` / ``start_conversation()`` — entry points
"""

from __future__ import annotations

import json
import time
from typing import Any
from uuid import UUID

import msgspec
from opentelemetry import trace
from opentelemetry.trace import Span, StatusCode

from .init import require_initialized
from .otel import (
    ATTR_AGENT,
    ATTR_CONTEXT_SYSTEM_PERSONA,
    ATTR_CONTEXT_SYSTEM_TOOLS,
    ATTR_CONVERSATION_ID,
    ATTR_CONVERSATION_MODEL,
    ATTR_CONVERSATION_TAGS,
    ATTR_TOOL_CALL_ID,
    ATTR_TOOL_ERROR,
    ATTR_TOOL_INPUT,
    ATTR_TOOL_NAME,
    ATTR_TOOL_OUTPUT,
    ATTR_TURN_ITEMS,
    ATTR_TURN_ROLE,
    ATTR_TURN_START_TIME,
    EVENT_COST,
    EVENT_LLM_USAGE,
    EVENT_TURN,
    OtelAttributes,
    cost_delta_to_otel,
    llm_usage_to_otel,
)
from .span import set_span_error

_TRACER_NAME = "yuutrace"


# ---------------------------------------------------------------------------
# JSON serialization helpers
# ---------------------------------------------------------------------------


def _json_default(obj: object) -> Any:
    """Fallback for ``json.dumps`` that handles msgspec Structs."""
    if hasattr(obj, "__struct_fields__"):
        return msgspec.to_builtins(obj)
    raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")


def _serialize_item(item: Any) -> Any:
    """Normalize a single content item to a JSON-friendly dict."""
    if isinstance(item, dict):
        return item
    if hasattr(item, "__struct_fields__"):
        return msgspec.to_builtins(item)
    return item


def _items_to_json(items: list[Any]) -> str:
    """Serialize a list of items to a JSON string."""
    return json.dumps(
        [_serialize_item(it) for it in items],
        ensure_ascii=False,
        default=_json_default,
    )


# ---------------------------------------------------------------------------
# TurnContext
# ---------------------------------------------------------------------------


class TurnContext:
    """A single conversational turn (user or assistant).

    Each turn is recorded as an independent child span (named ``"turn"``)
    so that ``SimpleSpanProcessor`` exports it immediately on ``end()``,
    enabling real-time visibility in the trace UI.

    Usage::

        # Manual lifecycle
        turn = chat.start_turn("assistant")
        turn.add(*items)
        turn.usage(usage_obj, cost=cost_obj)
        turn.end()

        # Context manager
        with chat.turn("user") as t:
            t.add(text_item, image_item)
    """

    __slots__ = ("_span", "_role", "_items", "_ended", "_start_ns")

    def __init__(self, parent_span: Span, role: str, conversation_id: str) -> None:
        tracer = trace.get_tracer(_TRACER_NAME)
        self._span = tracer.start_span(
            "turn",
            context=trace.set_span_in_context(parent_span),
        )
        self._span.set_attribute(ATTR_CONVERSATION_ID, conversation_id)
        self._span.set_attribute(ATTR_TURN_ROLE, role)
        self._role = role
        self._items: list[Any] = []
        self._ended = False
        self._start_ns = time.time_ns()

    def add(self, *items: Any) -> None:
        """Append content items (yuullm.Item dicts or msgspec Structs)."""
        self._items.extend(items)

    def usage(self, u: object, cost: object | None = None) -> None:
        """Record token usage (and optional cost) for this turn.

        Emits ``yuu.llm.usage`` (and optionally ``yuu.cost``) events on the
        turn's own child span.  Also stores the attributes directly on the span
        so the frontend can render per-turn metrics without parsing events.
        """
        from .usage import _to_llm_usage_delta

        from .types import CostCategory, CostDelta, Currency, LlmUsageDelta

        if not isinstance(u, LlmUsageDelta):
            u = _to_llm_usage_delta(u)

        # Emit usage event on the turn span
        usage_otel = llm_usage_to_otel(u)
        self._span.add_event(EVENT_LLM_USAGE, attributes=usage_otel)  # type: ignore[arg-type]
        # Also set as span attributes for direct access
        for k, v in usage_otel.items():
            self._span.set_attribute(k, v)

        # Emit cost event if provided
        if cost is not None:
            cost_delta = CostDelta(
                category=CostCategory.llm,
                currency=Currency.USD,
                amount=cost.total_cost,  # type: ignore[attr-defined]
                source=getattr(cost, "source", None),
                llm_provider=u.provider,
                llm_model=u.model,
                llm_request_id=u.request_id,
            )
            cost_otel = cost_delta_to_otel(cost_delta)
            self._span.add_event(EVENT_COST, attributes=cost_otel)  # type: ignore[arg-type]
            for k, v in cost_otel.items():
                self._span.set_attribute(k, v)

    def end(self, error: Exception | None = None) -> None:
        """End the turn's child span, triggering immediate export."""
        if self._ended:
            return
        self._ended = True

        self._span.set_attribute(ATTR_TURN_ITEMS, _items_to_json(self._items))
        self._span.set_attribute(ATTR_TURN_START_TIME, self._start_ns)

        if error is not None:
            set_span_error(self._span, error)

        self._span.end()

    def __enter__(self) -> TurnContext:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: object,
    ) -> bool:
        self.end(error=exc_val if isinstance(exc_val, Exception) else None)
        return False


# ---------------------------------------------------------------------------
# ToolSpan
# ---------------------------------------------------------------------------


class ToolSpan:
    """Wraps a child span for a single tool invocation.

    Usage::

        # Context manager
        with tools.tool(name="search", call_id="tc_1", input={"q": "x"}) as ts:
            ts.ok("result")

        # Manual
        ts = tools.start_tool(name="search", call_id="tc_1", input={"q": "x"})
        ts.ok("result")
        ts.end()
    """

    __slots__ = ("_span", "_ended")

    def __init__(self, span: Span) -> None:
        self._span = span
        self._ended = False

    def ok(self, result: Any) -> None:
        """Record a successful tool result."""
        output = result if isinstance(result, str) else json.dumps(result, default=_json_default)
        self._span.set_attribute(ATTR_TOOL_OUTPUT, output)

    def fail(self, error: str) -> None:
        """Record a tool error."""
        self._span.set_attribute(ATTR_TOOL_ERROR, error)
        self._span.set_status(StatusCode.ERROR, error)

    def end(self) -> None:
        """End the tool span."""
        if self._ended:
            return
        self._ended = True
        self._span.end()

    def __enter__(self) -> ToolSpan:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: object,
    ) -> bool:
        if exc_val is not None:
            self.fail(f"{type(exc_val).__name__}: {exc_val}")
        self.end()
        return False


# ---------------------------------------------------------------------------
# ToolsContext
# ---------------------------------------------------------------------------


class ToolsContext:
    """Wraps a child span for a batch of tool executions.

    Usage::

        # Context manager
        with chat.tools() as tools:
            with tools.tool(name="search", call_id="tc_1", input={}) as ts:
                ts.ok("ok")

        # Manual
        tools = chat.start_tools()
        ts = tools.start_tool(name="search", call_id="tc_1", input={})
        ts.ok("ok")
        ts.end()
        tools.end()
    """

    __slots__ = ("_span", "_conversation_id", "_ended")

    def __init__(self, span: Span, conversation_id: str) -> None:
        self._span = span
        self._conversation_id = conversation_id
        self._ended = False

    def start_tool(
        self,
        *,
        name: str,
        call_id: str,
        input: Any,  # noqa: A002
    ) -> ToolSpan:
        """Start a child span for a single tool invocation."""
        tracer = trace.get_tracer(_TRACER_NAME)
        tool_span = tracer.start_span(
            f"tool:{name}",
            context=trace.set_span_in_context(self._span),
        )
        tool_span.set_attribute(ATTR_TOOL_NAME, name)
        tool_span.set_attribute(ATTR_TOOL_CALL_ID, call_id)
        tool_span.set_attribute(ATTR_CONVERSATION_ID, self._conversation_id)
        input_str = input if isinstance(input, str) else json.dumps(input, default=_json_default)
        tool_span.set_attribute(ATTR_TOOL_INPUT, input_str)
        return ToolSpan(tool_span)

    def tool(
        self,
        *,
        name: str,
        call_id: str,
        input: Any,  # noqa: A002
    ) -> ToolSpan:
        """Context manager shortcut for ``start_tool``."""
        return self.start_tool(name=name, call_id=call_id, input=input)

    def end(self) -> None:
        """End the tools batch span."""
        if self._ended:
            return
        self._ended = True
        self._span.end()

    def __enter__(self) -> ToolsContext:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: object,
    ) -> bool:
        self.end()
        return False


# ---------------------------------------------------------------------------
# ConversationContext
# ---------------------------------------------------------------------------


class ConversationContext:
    """Wraps a conversation span. Entry point for recording turns and tools.

    Usage::

        with ytrace.conversation(id=uuid, agent="a", model="m") as chat:
            chat.system("You are helpful.", tools=[...])
            chat.user(text_item, image_item)

            with chat.turn("assistant") as t:
                t.add(*response_items)
                t.usage(usage)

            with chat.tools() as batch:
                with batch.tool(name="search", call_id="tc_1", input={}) as ts:
                    ts.ok("result")
    """

    __slots__ = ("_span", "_conversation_id", "_ended")

    def __init__(self, span: Span, conversation_id: str) -> None:
        self._span = span
        self._conversation_id = conversation_id
        self._ended = False

    # -- Turn API -----------------------------------------------------------

    def start_turn(self, role: str) -> TurnContext:
        """Start a turn manually. Caller must call ``turn.end()``."""
        return TurnContext(self._span, role, self._conversation_id)

    def turn(self, role: str) -> TurnContext:
        """Context manager for a turn."""
        return TurnContext(self._span, role, self._conversation_id)

    def user(self, *items: Any) -> None:
        """Sugar: record an instant user turn.

        Accepts ``yuullm.Item`` dicts, ``msgspec.Struct`` instances, or plain
        ``str`` (auto-wrapped to ``TextItem``).
        """
        normalized = [
            {"type": "text", "text": it} if isinstance(it, str) else it
            for it in items
        ]
        t = self.start_turn("user")
        t.add(*normalized)
        t.end()

    # -- System prompt ------------------------------------------------------

    def system(self, persona: str, tools: Any | None = None) -> None:
        """Record the system prompt on the conversation span."""
        self._span.set_attribute(ATTR_CONTEXT_SYSTEM_PERSONA, persona)
        if tools is not None:
            tools_str = tools if isinstance(tools, str) else json.dumps(tools, default=_json_default)
            self._span.set_attribute(ATTR_CONTEXT_SYSTEM_TOOLS, tools_str)

    # -- Tool batch ---------------------------------------------------------

    def start_tool_batch(self) -> ToolsContext:
        """Start a tool execution batch span."""
        tracer = trace.get_tracer(_TRACER_NAME)
        tools_span = tracer.start_span(
            "tools",
            context=trace.set_span_in_context(self._span),
        )
        tools_span.set_attribute(ATTR_CONVERSATION_ID, self._conversation_id)
        return ToolsContext(tools_span, self._conversation_id)

    def start_tools(self) -> ToolsContext:
        """Alias for ``start_tool_batch()``."""
        return self.start_tool_batch()

    def tools(self) -> ToolsContext:
        """Context manager version of ``start_tool_batch()``."""
        return self.start_tool_batch()

    # -- Lifecycle ----------------------------------------------------------

    def end(self, error: Exception | None = None) -> None:
        """End the conversation span."""
        if self._ended:
            return
        self._ended = True
        if error is not None:
            set_span_error(self._span, error)
        self._span.end()

    def __enter__(self) -> ConversationContext:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: object,
    ) -> bool:
        self.end(error=exc_val if isinstance(exc_val, Exception) else None)
        return False


# ---------------------------------------------------------------------------
# Module-level entry points
# ---------------------------------------------------------------------------


def start_conversation(
    *,
    id: UUID,  # noqa: A002
    agent: str,
    model: str,
    tags: list[str] | None = None,
) -> ConversationContext:
    """Start a conversation span (manual lifecycle).

    Caller must call ``.end()`` when the conversation is done.
    """
    require_initialized()
    conversation_id = str(id)

    tracer = trace.get_tracer(_TRACER_NAME)
    span = tracer.start_span("conversation")
    span.set_attribute(ATTR_CONVERSATION_ID, conversation_id)
    span.set_attribute(ATTR_AGENT, agent)
    span.set_attribute(ATTR_CONVERSATION_MODEL, model)
    if tags:
        span.set_attribute(ATTR_CONVERSATION_TAGS, tags)
    return ConversationContext(span, conversation_id)


def conversation(
    *,
    id: UUID,  # noqa: A002
    agent: str,
    model: str,
    tags: list[str] | None = None,
) -> ConversationContext:
    """Context manager that starts and ends a conversation span.

    Usage::

        with ytrace.conversation(id=uuid, agent="a", model="m") as chat:
            chat.user(text_item)
            ...
    """
    return start_conversation(id=id, agent=agent, model=model, tags=tags)
