from __future__ import annotations

import json
import uuid
from dataclasses import dataclass, field
from typing import Any, Callable, TypeVar

from opentelemetry import context, trace
from opentelemetry.context import Context
from opentelemetry.trace import Span, StatusCode
from .debug_mode import _lemma_debug
from .experiment_mode import is_experiment_mode_enabled

T = TypeVar("T")
Input = TypeVar("Input")


@dataclass
class TraceContext:
    """Context object passed to the wrapped agent function."""

    span: Span
    """The active OpenTelemetry span for this agent run."""

    run_id: str
    """Unique identifier for this agent run."""
    auto_end_root: bool = True
    _root_ended: bool = field(default=False, init=False, repr=False)

    def on_complete(self, result: Any) -> bool:
        """Signal the run is complete.

        When ``auto_end_root`` is enabled (the default), this is a no-op and
        returns ``False`` — the span is ended automatically when the wrapped
        function returns.  When ``auto_end_root`` is disabled, this ends the
        top-level span and returns ``True``.
        """
        if self.auto_end_root or self._root_ended:
            _lemma_debug("trace-wrapper", "on_complete called but span not ended (auto_end_root active or already ended)", run_id=self.run_id)
            return False
        self.span.set_attribute("ai.agent.output", json.dumps(result, default=str))
        self._root_ended = True
        self.span.end()
        _lemma_debug("trace-wrapper", "span ended via on_complete", run_id=self.run_id)
        return True

    def record_error(self, error: Any) -> None:
        """Record an error on the span. Marks the span as errored."""
        exc = error if isinstance(error, BaseException) else Exception(str(error))
        self.span.record_exception(exc)
        self.span.set_status(StatusCode.ERROR)


def wrap_agent(
    agent_name: str,
    fn: Callable[[TraceContext, Input], T],
    *,
    is_experiment: bool = False,
    auto_end_root: bool = True,
) -> Callable[[Input], tuple[T, str, Span]]:
    """Wrap an agent function with OpenTelemetry tracing.

    Creates a new span on every invocation, attaches agent metadata
    (run ID, experiment flag), and handles error recording.

    Sets ``ai.agent.input`` and ``ai.agent.output`` as JSON strings for Lemma ingestion.

    Args:
        agent_name: Human-readable name used as the span name.
        fn: The agent function to wrap. Receives a :class:`TraceContext`
            as its first argument and the call-time ``input`` as its second.
        is_experiment: Mark this run as an experiment in Lemma.
        auto_end_root: Automatically end the top-level span when the wrapped
            function returns or throws (default: ``True``).  Set to ``False``
            to manage span lifetime manually via :meth:`TraceContext.on_complete`.

    Returns:
        A wrapper that accepts an ``input``, calls *fn* inside a traced
        context, and returns ``(result, run_id, span)``.

    Example::

        from typing import TypedDict

        class AgentInput(TypedDict):
            topic: str

        async def handler(ctx: TraceContext, input: AgentInput) -> str:
            result = await do_work(input["topic"])
            return result

        my_agent = wrap_agent("my-agent", handler)
        await my_agent({"topic": "math"})
    """

    def _start_root_span(agent_input: Input) -> tuple[Span, str]:
        tracer = trace.get_tracer("lemma")
        run_id = str(uuid.uuid4())

        span = tracer.start_span(
            "ai.agent.run",
            context=Context(),
            attributes={
                "ai.agent.name": agent_name,
                "lemma.run_id": run_id,
                "lemma.is_experiment": is_experiment_mode_enabled() or is_experiment,
                "lemma.auto_end_root": auto_end_root,
            },
        )
        span.set_attribute("ai.agent.input", json.dumps(agent_input, default=str))
        _lemma_debug("trace-wrapper", "span started", agent_name=agent_name, run_id=run_id, auto_end_root=auto_end_root)
        return span, run_id

    async def _wrapped_async(input: Input) -> tuple[T, str, Span]:
        import asyncio  # noqa: F811 – deferred so sync callers don't pay the import

        span, run_id = _start_root_span(input)

        ctx = trace.set_span_in_context(span, Context())
        token = context.attach(ctx)

        try:
            trace_ctx = TraceContext(
                span=span,
                run_id=run_id,
                auto_end_root=auto_end_root,
            )

            if asyncio.iscoroutinefunction(fn):
                result = await fn(trace_ctx, input)
            else:
                result = fn(trace_ctx, input)  # pragma: no cover

            if auto_end_root and not trace_ctx._root_ended:
                span.set_attribute("ai.agent.output", json.dumps(result, default=str))
                trace_ctx._root_ended = True
                span.end()
                _lemma_debug("trace-wrapper", "span auto-ended after fn returned", run_id=run_id)

            return result, run_id, span
        except BaseException as exc:
            span.record_exception(exc)
            span.set_status(StatusCode.ERROR)
            if not trace_ctx._root_ended:
                trace_ctx._root_ended = True
                span.end()
                _lemma_debug("trace-wrapper", "span ended on error", run_id=run_id, error=str(exc))
            raise
        finally:
            context.detach(token)

    def _wrapped_sync(input: Input) -> tuple[T, str, Span]:
        span, run_id = _start_root_span(input)

        ctx = trace.set_span_in_context(span, Context())
        token = context.attach(ctx)

        try:
            trace_ctx = TraceContext(
                span=span,
                run_id=run_id,
                auto_end_root=auto_end_root,
            )
            result = fn(trace_ctx, input)

            if auto_end_root and not trace_ctx._root_ended:
                span.set_attribute("ai.agent.output", json.dumps(result, default=str))
                trace_ctx._root_ended = True
                span.end()
                _lemma_debug("trace-wrapper", "span auto-ended after fn returned", run_id=run_id)

            return result, run_id, span
        except BaseException as exc:
            span.record_exception(exc)
            span.set_status(StatusCode.ERROR)
            if not trace_ctx._root_ended:
                trace_ctx._root_ended = True
                span.end()
                _lemma_debug("trace-wrapper", "span ended on error", run_id=run_id, error=str(exc))
            raise
        finally:
            context.detach(token)

    import asyncio

    if asyncio.iscoroutinefunction(fn):
        return _wrapped_async  # type: ignore[return-value]
    return _wrapped_sync  # type: ignore[return-value]


@dataclass
class RunContext:
    """Context object yielded by the :class:`lemma_run` context manager.

    Unlike :class:`TraceContext`, this is not associated with ``auto_end_root`` —
    the span always ends when the ``with`` / ``async with`` block exits.
    """

    span: Span
    """The active OpenTelemetry span for this agent run."""

    run_id: str
    """Unique identifier for this agent run."""

    _ended: bool = field(default=False, init=False, repr=False)

    def on_complete(self, result: Any) -> None:
        """Record the run output.

        Sets ``ai.agent.output`` on the span. The span itself is ended by the
        context manager when the ``with`` block exits, not by this call.
        """
        self.span.set_attribute("ai.agent.output", json.dumps(result, default=str))
        _lemma_debug("trace-wrapper", "on_complete called (context manager)", run_id=self.run_id)

    def record_error(self, error: Any) -> None:
        """Record an error on the span. Marks the span as errored."""
        exc = error if isinstance(error, BaseException) else Exception(str(error))
        self.span.record_exception(exc)
        self.span.set_status(StatusCode.ERROR)


class lemma_run:
    """Context manager that opens a Lemma agent run span for the duration of the block.

    Use this as a zero-refactor alternative to :func:`wrap_agent` when you want
    to instrument existing code without extracting it into a separate function.

    Supports both ``async with`` and ``with``.

    Args:
        agent_name: Human-readable name recorded as ``ai.agent.name``.
        input: The run input. Serialized as ``ai.agent.input`` JSON attribute.
        is_experiment: Mark this run as an experiment in Lemma.

    Example (async)::

        async with lemma_run("my-agent", input=user_message) as run:
            response = await call_llm(user_message)
            run.on_complete(response)

        print(run.run_id)

    Example (sync)::

        with lemma_run("my-agent", input=user_message) as run:
            response = call_llm(user_message)
            run.on_complete(response)
    """

    def __init__(
        self,
        agent_name: str,
        *,
        input: Any = None,
        is_experiment: bool = False,
    ) -> None:
        self._agent_name = agent_name
        self._input = input
        self._is_experiment = is_experiment
        self._token: object | None = None
        self._run: RunContext | None = None

    def _start(self) -> RunContext:
        tracer = trace.get_tracer("lemma")
        run_id = str(uuid.uuid4())

        span = tracer.start_span(
            "ai.agent.run",
            context=Context(),
            attributes={
                "ai.agent.name": self._agent_name,
                "lemma.run_id": run_id,
                "lemma.is_experiment": is_experiment_mode_enabled() or self._is_experiment,
                "lemma.auto_end_root": False,
            },
        )
        span.set_attribute("ai.agent.input", json.dumps(self._input, default=str))
        _lemma_debug("trace-wrapper", "span started (context manager)", agent_name=self._agent_name, run_id=run_id)

        ctx = trace.set_span_in_context(span, Context())
        self._token = context.attach(ctx)
        self._run = RunContext(span=span, run_id=run_id)
        return self._run

    def _end(self, exc_val: BaseException | None) -> None:
        run = self._run
        if run is None:
            return

        if exc_val is not None:
            run.record_error(exc_val)

        if not run._ended:
            run._ended = True
            run.span.end()
            _lemma_debug("trace-wrapper", "span ended (context manager exit)", run_id=run.run_id)

        if self._token is not None:
            context.detach(self._token)
            self._token = None

    # ---- sync ----------------------------------------------------------------

    def __enter__(self) -> RunContext:
        return self._start()

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: object,
    ) -> bool:
        self._end(exc_val)
        return False  # never suppress exceptions

    # ---- async ---------------------------------------------------------------

    async def __aenter__(self) -> RunContext:
        return self._start()

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: object,
    ) -> bool:
        self._end(exc_val)
        return False  # never suppress exceptions
