"""
The central coordinator for all tracing operations in agenx

Tracer has three responsibilities:
  - Span construction: builds Span + LiveSpan with correct IDs, timestamps, and attribute class for the requested type
  - Context wiring: reads the active span from ContextVar to auto-set parent_id and trace_id. Never pass these manually
  - Exporter management — holds the callback list that fires when spans finish
"""

from __future__ import annotations

from contextlib import asynccontextmanager, contextmanager
from typing import Callable

from agenx._internal import (
    now,
    pop_span,
    push_span,
    get_current_span,
    make_span_id,
    make_trace_id
)
from agenx.trace.spans.models import SPAN_TYPE_TO_ATTRIBUTES, Span, SpanType
from agenx.trace.spans.span import LiveSpan
from agenx.decorators import create_decorator


class Tracer:
    def __init__(self) -> None:
        self._exporters: list[Callable] = []

    def configure(self, exporters: list[Callable] | None = None) -> None:
        self._exporters = list(exporters) if exporters else []

    def _create_span(
        self,
        name: str,
        *,
        type: SpanType = SpanType.STEP,
        trace_id: str | None = None,
        parent_id: str | None = None,
    ) -> LiveSpan:
        current = get_current_span()

        resolved_parent_id = parent_id or (current.span_id if current else None)
        resolved_trace_id = trace_id or (current.trace_id if current else make_trace_id())

        attributes_cls = SPAN_TYPE_TO_ATTRIBUTES[type]

        span = Span(
            span_id=make_span_id(),
            trace_id=resolved_trace_id,
            parent_id=resolved_parent_id,
            name=name,
            type=type,
            start_time=now(),
            attributes=attributes_cls()
        )

        return LiveSpan(span, on_end_callbacks=list(self._exporters))
    
    @contextmanager
    def span(
        self,
        name: str,
        *,
        type: SpanType = SpanType.STEP,
        trace_id: str | None = None,
        parent_id: str | None = None
    ):
        live = self._create_span(
            name,
            type=type,
            trace_id=trace_id,
            parent_id=parent_id
        )

        token = push_span(live)

        try:
            yield live
        except BaseException as exc:
            live._finish(exc=exc)
            raise
        else:
            live._finish()
        finally:
            pop_span(token)
    
    @asynccontextmanager
    async def aspan(
        self,
        name: str,
        *,
        type: SpanType = SpanType.STEP,
        trace_id: str | None = None,
        parent_id: str | None = None
    ):
        live = self._create_span(
            name, 
            type=type, 
            trace_id=trace_id, 
            parent_id=parent_id
        )

        token = push_span(live)

        try:
            yield live
        except BaseException as exc:
            live._finish(exc=exc)
            raise
        else:
            live._finish()
        finally:
            pop_span(token)

    def step(self, func=None, /, name=None, capture_args=True, capture_result=True):
        """Trace a generic step or function"""
        return create_decorator(
            func=func,
            name=name,
            span_type=SpanType.STEP,
            tracer=self,
            capture_args=capture_args,
            capture_result=capture_result,
        )
    
    def llm(self, func=None, /, name=None, model="", provider="", temperature=0.0, capture_args=True, capture_result=True):
        """Trace an LLM call"""
        return create_decorator(
            func=func,
            name=name,
            span_type=SpanType.LLM,
            tracer=self,
            model=model,
            provider=provider,
            temperature=temperature,
            capture_args=capture_args,
            capture_result=capture_result,
        )

    def agent(self, func=None, /, name=None, tools=None, capture_args=True, capture_result=True):
        """Trace an agent workflow"""
        return create_decorator(
            func=func,
            name=name,
            span_type=SpanType.AGENT,
            tracer=self,
            tools=tools or [],
            capture_args=capture_args,
            capture_result=capture_result,
        )
    
    def tool(self, func=None, /, name=None, tool_name="", capture_args=True, capture_result=True):
        """Trace a tool call."""
        actual_tool_name = tool_name if tool_name else (func.__name__ if func else "")
        return create_decorator(
            func=func,
            name=name,
            span_type=SpanType.TOOL,
            tracer=self,
            tool_name=actual_tool_name,
            capture_args=capture_args,
            capture_result=capture_result,
        )

    def retriever(self, func=None, /, name=None, top_k=0, capture_args=True, capture_result=True):
        """Trace a document retrieval call."""
        return create_decorator(
            func=func,
            name=name,
            span_type=SpanType.RETRIEVER,
            tracer=self,
            top_k=top_k,
            capture_args=capture_args,
            capture_result=capture_result,
        )
    
    def guard(self, func=None, /, name=None, guard_name="", capture_args=True, capture_result=True):
        """Trace a guardrail/validation check."""
        actual_guard_name = guard_name if guard_name else (func.__name__ if func else "")
        return create_decorator(
            func=func,
            name=name,
            span_type=SpanType.GUARD,
            tracer=self,
            guard_name=actual_guard_name,
            capture_args=capture_args,
            capture_result=capture_result,
        )
