"""W3C trace context propagation via contextvars."""
from __future__ import annotations

import contextlib
import logging
from contextvars import ContextVar
from datetime import datetime, timezone
from typing import Any, Generator
from uuid import uuid4

from fibinaci.tracing.span_builder import SpanData

logger = logging.getLogger("fibinaci.tracing")

# ContextVars for trace propagation
_current_trace_id: ContextVar[str | None] = ContextVar("fibinaci_trace_id", default=None)
_current_span_id: ContextVar[str | None] = ContextVar("fibinaci_span_id", default=None)
_current_spans: ContextVar[list[SpanData]] = ContextVar("fibinaci_spans")


def _new_trace_id() -> str:
    return uuid4().hex


def _new_span_id() -> str:
    return uuid4().hex[:16]


@contextlib.contextmanager
def trace(
    name: str | None = None,
    *,
    session_id: str | None = None,
    user_id: str | None = None,
    workflow_name: str | None = None,
    **kwargs: Any,
) -> Generator[str, None, None]:
    """Start a new trace context. Yields the trace_id."""
    trace_id = _new_trace_id()
    token_tid = _current_trace_id.set(trace_id)
    token_spans = _current_spans.set([])

    try:
        yield trace_id
    finally:
        # Flush all spans from this trace
        spans = _current_spans.get()
        if spans:
            try:
                from fibinaci._globals import get_buffer
                buf = get_buffer()
                for sd in spans:
                    sd.session_id = session_id or sd.session_id
                    sd.user_id = user_id or sd.user_id
                    sd.workflow_name = workflow_name or sd.workflow_name
                    buf.push(sd.to_dict())
            except Exception:
                logger.debug("Failed to flush trace spans", exc_info=True)

        _current_trace_id.reset(token_tid)
        _current_spans.reset(token_spans)


@contextlib.contextmanager
def span(
    name: str | None = None,
    *,
    span_kind: str = "custom",
    provider: str | None = None,
    model_type: str | None = None,
    **kwargs: Any,
) -> Generator[SpanData, None, None]:
    """Create a span within the current trace. Yields SpanData for enrichment."""
    trace_id = _current_trace_id.get()
    if trace_id is None:
        # No active trace — create a standalone span
        trace_id = _new_trace_id()

    parent_span_id = _current_span_id.get()
    span_id = _new_span_id()
    token = _current_span_id.set(span_id)

    sd = SpanData(
        trace_id=trace_id,
        span_id=span_id,
        parent_span_id=parent_span_id,
        span_kind=span_kind,
        provider=provider,
        model_type=model_type,
        started_at=datetime.now(timezone.utc),
        **kwargs,
    )

    try:
        yield sd
        sd.status = sd.status or "ok"
    except Exception as exc:
        sd.status = "error"
        sd.error_type = type(exc).__name__
        sd.error_message = str(exc)
        raise
    finally:
        sd.ended_at = sd.ended_at or datetime.now(timezone.utc)
        sd.duration_ms = (sd.ended_at - sd.started_at).total_seconds() * 1000

        # Attach to current trace's span list (flushed when trace exits)
        try:
            spans = _current_spans.get()
            spans.append(sd)
        except LookupError:
            # span() called outside of trace() — flush immediately
            from fibinaci._globals import get_buffer
            buf = get_buffer()
            if buf is not None:
                buf.push(sd.to_dict())
            pass

        _current_span_id.reset(token)


def outcome(
    *,
    trace_id: str | None = None,
    success: bool,
    score: float | None = None,
    **kwargs: Any,
) -> None:
    """Record an outcome for the current or specified trace."""
    tid = trace_id or _current_trace_id.get()
    if tid is None:
        logger.warning("outcome() called outside of a trace context and no trace_id provided")
        return

    try:
        from fibinaci._globals import get_buffer
        buf = get_buffer()
        buf.push({
            "_fibinaci_type": "outcome",
            "trace_id": tid,
            "success": success,
            "score": score,
            **kwargs,
        })
    except Exception:
        logger.debug("Failed to record outcome", exc_info=True)


def get_current_trace_id() -> str | None:
    return _current_trace_id.get()


def get_current_span_id() -> str | None:
    return _current_span_id.get()
