"""Tracing context managers and decorator for grouping LLM calls into traces and spans."""

from __future__ import annotations

import asyncio
import functools
import inspect
import json
import logging
import time
import uuid
from contextvars import ContextVar
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Callable, TypeVar, overload

from llm_observatory.config import get_config
from llm_observatory.transport import get_transport

logger = logging.getLogger("llm_observatory")


@dataclass
class _SpanContext:
    trace_id: str
    span_id: str
    parent_id: str | None
    span_type: str  # "trace", "span"
    name: str
    start_time: float  # time.time()
    input: Any = None
    output: Any = None
    status: str = "ok"
    metadata: dict[str, Any] = field(default_factory=dict)


_current_span: ContextVar[_SpanContext | None] = ContextVar("_current_span", default=None)


def get_current_trace_id() -> str | None:
    ctx = _current_span.get()
    return ctx.trace_id if ctx else None


def get_current_span_id() -> str | None:
    ctx = _current_span.get()
    return ctx.span_id if ctx else None


def get_current_traceparent() -> str | None:
    """Return W3C-style traceparent string for cross-service propagation.

    Format: "00-{trace_id}-{span_id}-01" where IDs have hyphens stripped.
    """
    ctx = _current_span.get()
    if ctx is None:
        return None
    # Strip hyphens from UUIDs to produce W3C-compliant hex IDs
    tid = ctx.trace_id.replace("-", "")
    sid = ctx.span_id.replace("-", "")
    return f"00-{tid}-{sid}-01"


def _parse_traceparent(traceparent: str) -> tuple[str, str] | None:
    """Parse W3C-style traceparent into (trace_id, parent_span_id).

    Expected format: "00-{hex_trace_id}-{hex_span_id}-01"
    Returns None on invalid input.
    """
    parts = traceparent.split("-")
    if len(parts) != 4:
        return None
    trace_id, span_id = parts[1], parts[2]
    try:
        int(trace_id, 16)
        int(span_id, 16)
    except ValueError:
        return None
    return trace_id, span_id


def _send_span_event(ctx: _SpanContext) -> None:
    """Build and enqueue a trace/span event."""
    try:
        cfg = get_config()
        if not cfg.get("enabled"):
            return

        event: dict[str, Any] = {
            "event_id": ctx.span_id,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "trace_id": ctx.trace_id,
            "parent_id": ctx.parent_id,
            "span_type": ctx.span_type,
            "span_name": ctx.name,
            "span_start_time": datetime.fromtimestamp(ctx.start_time, tz=timezone.utc).isoformat(),
            "span_end_time": datetime.now(timezone.utc).isoformat(),
            "span_status": ctx.status,
            "request": {"model": "", "messages": [], "stream": False},
            "response": {},
            "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
            "metrics": {"latency_ms": int((time.time() - ctx.start_time) * 1000)},
            "metadata": {},
        }

        if ctx.input is not None:
            event["span_input"] = json.dumps(ctx.input) if not isinstance(ctx.input, str) else ctx.input
        if ctx.output is not None:
            event["span_output"] = json.dumps(ctx.output) if not isinstance(ctx.output, str) else ctx.output

        get_transport().enqueue(event)
    except Exception as exc:
        logger.debug("llm_observatory: failed to send span event: %s", exc)


class trace:
    """Context manager that creates a trace grouping related LLM calls and spans."""

    def __init__(
        self,
        name: str = "",
        input: Any = None,
        traceparent: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        self._name = name
        self._input = input
        self._traceparent = traceparent
        self._metadata = metadata or {}
        self._ctx: _SpanContext | None = None
        self._token: Any = None

    def __enter__(self) -> trace:
        try:
            trace_id = uuid.uuid4().hex
            parent_id: str | None = None

            if self._traceparent:
                parsed = _parse_traceparent(self._traceparent)
                if parsed:
                    trace_id = parsed[0]
                    parent_id = parsed[1]

            self._ctx = _SpanContext(
                trace_id=trace_id,
                span_id=uuid.uuid4().hex,
                parent_id=parent_id,
                span_type="trace",
                name=self._name,
                start_time=time.time(),
                input=self._input,
                metadata=self._metadata,
            )
            self._token = _current_span.set(self._ctx)
        except Exception as exc:
            logger.debug("llm_observatory: failed to enter trace context: %s", exc)
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        try:
            if self._ctx is not None:
                if exc_type is not None:
                    self._ctx.status = "error"
                _send_span_event(self._ctx)
            if self._token is not None:
                _current_span.reset(self._token)
        except Exception as exc:
            logger.debug("llm_observatory: failed to exit trace context: %s", exc)
        return None  # Don't suppress exceptions

    def set_output(self, value: Any) -> None:
        if self._ctx is not None:
            self._ctx.output = value

    def update(self, name: str | None = None, metadata: dict[str, Any] | None = None) -> None:
        if self._ctx is not None:
            if name is not None:
                self._ctx.name = name
            if metadata is not None:
                self._ctx.metadata.update(metadata)


class span:
    """Context manager that creates a span within the current trace."""

    def __init__(
        self,
        name: str = "",
        input: Any = None,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        self._name = name
        self._input = input
        self._metadata = metadata or {}
        self._ctx: _SpanContext | None = None
        self._token: Any = None
        self._parent_ctx: _SpanContext | None = None

    def __enter__(self) -> span:
        try:
            self._parent_ctx = _current_span.get()

            if self._parent_ctx is not None:
                trace_id = self._parent_ctx.trace_id
                parent_id = self._parent_ctx.span_id
            else:
                # Span outside a trace: create an implicit trace_id
                trace_id = uuid.uuid4().hex
                parent_id = None

            self._ctx = _SpanContext(
                trace_id=trace_id,
                span_id=uuid.uuid4().hex,
                parent_id=parent_id,
                span_type="span",
                name=self._name,
                start_time=time.time(),
                input=self._input,
                metadata=self._metadata,
            )
            self._token = _current_span.set(self._ctx)
        except Exception as exc:
            logger.debug("llm_observatory: failed to enter span context: %s", exc)
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        try:
            if self._ctx is not None:
                if exc_type is not None:
                    self._ctx.status = "error"
                _send_span_event(self._ctx)
            if self._token is not None:
                _current_span.reset(self._token)
        except Exception as exc:
            logger.debug("llm_observatory: failed to exit span context: %s", exc)
        return None  # Don't suppress exceptions

    def set_output(self, value: Any) -> None:
        if self._ctx is not None:
            self._ctx.output = value

    def update(self, name: str | None = None, metadata: dict[str, Any] | None = None) -> None:
        if self._ctx is not None:
            if name is not None:
                self._ctx.name = name
            if metadata is not None:
                self._ctx.metadata.update(metadata)


F = TypeVar("F", bound=Callable[..., Any])


@overload
def observe(fn: F) -> F: ...


@overload
def observe(*, name: str = ..., capture_input: bool = ..., capture_output: bool = ...) -> Callable[[F], F]: ...


def observe(
    fn: F | None = None,
    *,
    name: str | None = None,
    capture_input: bool = True,
    capture_output: bool = True,
) -> F | Callable[[F], F]:
    """Decorator that wraps a function in a span, auto-capturing input/output.

    Usage:
        @observe
        def my_function(query: str): ...

        @observe(name="custom-name", capture_input=False)
        def my_function(query: str): ...

    The first call in a call stack becomes a trace; nested calls become spans.
    Works with both sync and async functions.
    """

    def decorator(func: F) -> F:
        span_name = name or func.__name__

        @functools.wraps(func)
        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
            captured_input = _capture_args(func, args, kwargs) if capture_input else None
            ctx_mgr = _pick_context_manager(span_name, captured_input)

            with ctx_mgr as s:
                try:
                    result = await func(*args, **kwargs)
                except Exception:
                    raise
                else:
                    if capture_output and result is not None:
                        s.set_output(result if isinstance(result, (str, int, float, bool, list, dict)) else str(result))
                    return result

        @functools.wraps(func)
        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
            captured_input = _capture_args(func, args, kwargs) if capture_input else None
            ctx_mgr = _pick_context_manager(span_name, captured_input)

            with ctx_mgr as s:
                try:
                    result = func(*args, **kwargs)
                except Exception:
                    raise
                else:
                    if capture_output and result is not None:
                        s.set_output(result if isinstance(result, (str, int, float, bool, list, dict)) else str(result))
                    return result

        if asyncio.iscoroutinefunction(func):
            return async_wrapper  # type: ignore[return-value]
        return sync_wrapper  # type: ignore[return-value]

    if fn is not None:
        return decorator(fn)
    return decorator  # type: ignore[return-value]


def _pick_context_manager(name: str, input_val: Any) -> trace | span:
    """Use trace if no active context, otherwise span."""
    if _current_span.get() is None:
        return trace(name=name, input=input_val)
    return span(name=name, input=input_val)


def _capture_args(func: Callable[..., Any], args: tuple, kwargs: dict) -> dict | None:
    """Build a dict of function arguments for span input."""
    try:
        sig = inspect.signature(func)
        bound = sig.bind(*args, **kwargs)
        bound.apply_defaults()
        captured = {}
        for k, v in bound.arguments.items():
            if isinstance(v, (str, int, float, bool, list, dict, type(None))):
                captured[k] = v
            else:
                captured[k] = str(v)
        return captured if captured else None
    except Exception as exc:
        logger.debug("llm_observatory: failed to capture args: %s", exc)
        return None


def emit_completed_span(
    name: str,
    traceparent: str,
    start_ns: int,
    end_ns: int,
) -> None:
    """Create a retroactive span with explicit start/end times.

    Used for measuring queue wait time: the producer records time.time_ns()
    when sending the message, and the consumer calls this to backfill the
    wait duration as a span in the trace.
    """
    try:
        cfg = get_config()
        if not cfg.get("enabled"):
            return

        parsed = _parse_traceparent(traceparent)
        if parsed is None:
            logger.debug("llm_observatory: invalid traceparent in emit_completed_span")
            return

        trace_id, parent_span_id = parsed
        start_time = start_ns / 1e9
        end_time = end_ns / 1e9
        latency_ms = int((end_time - start_time) * 1000)

        event: dict[str, Any] = {
            "event_id": uuid.uuid4().hex,
            "timestamp": datetime.fromtimestamp(end_time, tz=timezone.utc).isoformat(),
            "trace_id": trace_id,
            "parent_id": parent_span_id,
            "span_type": "span",
            "span_name": name,
            "span_start_time": datetime.fromtimestamp(start_time, tz=timezone.utc).isoformat(),
            "span_end_time": datetime.fromtimestamp(end_time, tz=timezone.utc).isoformat(),
            "span_status": "ok",
            "request": {"model": "", "messages": [], "stream": False},
            "response": {},
            "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
            "metrics": {"latency_ms": latency_ms},
            "metadata": {},
        }

        get_transport().enqueue(event)
    except Exception as exc:
        logger.debug("llm_observatory: failed to emit completed span: %s", exc)
