"""
核心追踪器

提供 Trace / Span 数据模型和 AgentTracer 追踪管理器。
"""

import logging
import time
import uuid
from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

logger = logging.getLogger(__name__)


class SpanType(Enum):
    """Span 类型"""

    LLM = "llm"  # LLM 调用
    TOOL = "tool"  # 工具执行
    TOOL_BATCH = "tool_batch"  # 工具批量执行
    MEMORY = "memory"  # 记忆操作
    CONTEXT = "context"  # 上下文管理（压缩等）
    REASONING = "reasoning"  # 推理循环
    PROMPT = "prompt"  # 提示词构建
    TASK = "task"  # 完整任务


class SpanStatus(Enum):
    """Span 状态"""

    OK = "ok"
    ERROR = "error"
    CANCELLED = "cancelled"


@dataclass
class Span:
    """
    单个操作的追踪记录。

    一个 Span 代表一次 LLM 调用、一次工具执行、一次记忆检索等。
    Span 可以嵌套形成父子关系。
    """

    span_id: str
    name: str
    span_type: SpanType
    start_time: float
    parent_id: str | None = None
    end_time: float | None = None
    status: SpanStatus = SpanStatus.OK
    attributes: dict[str, Any] = field(default_factory=dict)
    error_message: str | None = None

    @property
    def duration_ms(self) -> float | None:
        """耗时（毫秒）"""
        if self.end_time is None:
            return None
        return (self.end_time - self.start_time) * 1000

    def set_attribute(self, key: str, value: Any) -> None:
        """设置属性"""
        self.attributes[key] = value

    def set_error(self, message: str) -> None:
        """标记为错误"""
        self.status = SpanStatus.ERROR
        self.error_message = message

    def finish(self, status: SpanStatus | None = None) -> None:
        """结束 Span"""
        self.end_time = time.time()
        if status is not None:
            self.status = status

    def to_dict(self) -> dict[str, Any]:
        """序列化为字典"""
        result = {
            "span_id": self.span_id,
            "name": self.name,
            "type": self.span_type.value,
            "start_time": self.start_time,
            "status": self.status.value,
            "attributes": self.attributes,
        }
        if self.parent_id:
            result["parent_id"] = self.parent_id
        if self.end_time is not None:
            result["end_time"] = self.end_time
            result["duration_ms"] = self.duration_ms
        if self.error_message:
            result["error"] = self.error_message
        return result


@dataclass
class Trace:
    """
    一次完整的用户请求追踪。

    一个 Trace 包含多个 Span，代表从接收用户消息到返回响应的全过程。
    """

    trace_id: str
    session_id: str
    start_time: float
    spans: list[Span] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)
    end_time: float | None = None

    @property
    def duration_ms(self) -> float | None:
        """总耗时（毫秒）"""
        if self.end_time is None:
            return None
        return (self.end_time - self.start_time) * 1000

    @property
    def span_count(self) -> int:
        """Span 数量"""
        return len(self.spans)

    def add_span(self, span: Span) -> None:
        """添加 Span"""
        self.spans.append(span)

    def finish(self) -> None:
        """结束 Trace"""
        self.end_time = time.time()

    def get_summary(self) -> dict[str, Any]:
        """获取追踪摘要"""
        llm_spans = [s for s in self.spans if s.span_type == SpanType.LLM]
        tool_spans = [s for s in self.spans if s.span_type == SpanType.TOOL]

        total_input_tokens = sum(s.attributes.get("input_tokens", 0) for s in llm_spans)
        total_output_tokens = sum(s.attributes.get("output_tokens", 0) for s in llm_spans)

        tool_errors = sum(1 for s in tool_spans if s.status == SpanStatus.ERROR)

        return {
            "trace_id": self.trace_id,
            "session_id": self.session_id,
            "duration_ms": self.duration_ms,
            "total_spans": self.span_count,
            "llm_calls": len(llm_spans),
            "tool_calls": len(tool_spans),
            "tool_errors": tool_errors,
            "total_input_tokens": total_input_tokens,
            "total_output_tokens": total_output_tokens,
            "metadata": self.metadata,
        }

    def to_dict(self) -> dict[str, Any]:
        """序列化为字典"""
        return {
            "trace_id": self.trace_id,
            "session_id": self.session_id,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "duration_ms": self.duration_ms,
            "metadata": self.metadata,
            "spans": [s.to_dict() for s in self.spans],
            "summary": self.get_summary(),
        }


class AgentTracer:
    """
    Agent 追踪器。

    管理 Trace 和 Span 的生命周期，支持嵌套 Span 和多种导出器。

    Usage:
        tracer = AgentTracer()
        tracer.add_exporter(FileExporter("data/traces"))

        with tracer.start_trace("session-123") as trace:
            with tracer.llm_span(model="claude-4") as span:
                # ... LLM call ...
                span.set_attribute("input_tokens", 100)
    """

    def __init__(self, enabled: bool = True) -> None:
        self._enabled = enabled
        self._exporters: list[Any] = []  # TraceExporter instances
        self._current_trace: Trace | None = None
        self._span_stack: list[Span] = []

    @property
    def enabled(self) -> bool:
        return self._enabled

    def set_enabled(self, enabled: bool) -> None:
        self._enabled = enabled

    def add_exporter(self, exporter: Any) -> None:
        """添加追踪导出器"""
        self._exporters.append(exporter)

    @contextmanager
    def start_trace(self, session_id: str, **metadata: Any) -> Generator[Trace, None, None]:
        """
        开始一个新的 Trace。

        作为上下文管理器使用，退出时自动结束并导出。
        """
        if not self._enabled:
            # 返回一个空 Trace，不记录
            yield Trace(trace_id="", session_id=session_id, start_time=time.time())
            return

        trace = Trace(
            trace_id=str(uuid.uuid4()),
            session_id=session_id,
            start_time=time.time(),
            metadata=metadata,
        )
        self._current_trace = trace
        self._span_stack = []

        try:
            yield trace
        finally:
            trace.finish()
            self._export_trace(trace)
            self._current_trace = None
            self._span_stack = []

    def start_span(
        self,
        name: str,
        span_type: SpanType,
        parent: Span | None = None,
        **attributes: Any,
    ) -> Span:
        """
        创建并开始一个新的 Span。

        如果不指定 parent，自动使用 span stack 栈顶作为 parent。
        """
        if not self._enabled:
            return Span(span_id="", name=name, span_type=span_type, start_time=time.time())

        parent_id = None
        if parent:
            parent_id = parent.span_id
        elif self._span_stack:
            parent_id = self._span_stack[-1].span_id

        span = Span(
            span_id=str(uuid.uuid4()),
            name=name,
            span_type=span_type,
            start_time=time.time(),
            parent_id=parent_id,
            attributes=attributes,
        )

        if self._current_trace:
            self._current_trace.add_span(span)

        return span

    def end_span(self, span: Span, status: SpanStatus | None = None) -> None:
        """结束一个 Span"""
        if not self._enabled or not span.span_id:
            return
        span.finish(status)

    @contextmanager
    def span(
        self,
        name: str,
        span_type: SpanType,
        **attributes: Any,
    ) -> Generator[Span, None, None]:
        """
        通用 Span 上下文管理器。

        自动管理开始/结束和 span stack。
        """
        span = self.start_span(name, span_type, **attributes)
        self._span_stack.append(span)
        try:
            yield span
        except Exception as e:
            span.set_error(str(e))
            raise
        finally:
            self._span_stack.pop()
            span.finish()

    @contextmanager
    def llm_span(self, model: str = "", **attributes: Any) -> Generator[Span, None, None]:
        """LLM 调用 Span"""
        attrs = {"model": model, **attributes}
        with self.span("llm.call", SpanType.LLM, **attrs) as s:
            yield s

    @contextmanager
    def tool_span(self, tool_name: str = "", **attributes: Any) -> Generator[Span, None, None]:
        """工具执行 Span"""
        attrs = {"tool_name": tool_name, **attributes}
        with self.span("tool.execute", SpanType.TOOL, **attrs) as s:
            yield s

    @contextmanager
    def tool_batch_span(self, count: int = 0, **attributes: Any) -> Generator[Span, None, None]:
        """工具批量执行 Span"""
        attrs = {"tool_count": count, **attributes}
        with self.span("tool.batch", SpanType.TOOL_BATCH, **attrs) as s:
            yield s

    @contextmanager
    def memory_span(self, operation: str = "", **attributes: Any) -> Generator[Span, None, None]:
        """记忆操作 Span"""
        attrs = {"operation": operation, **attributes}
        with self.span("memory." + operation, SpanType.MEMORY, **attrs) as s:
            yield s

    @contextmanager
    def context_span(self, operation: str = "", **attributes: Any) -> Generator[Span, None, None]:
        """上下文操作 Span"""
        attrs = {"operation": operation, **attributes}
        with self.span("context." + operation, SpanType.CONTEXT, **attrs) as s:
            yield s

    @contextmanager
    def reasoning_span(self, iteration: int = 0, **attributes: Any) -> Generator[Span, None, None]:
        """推理循环 Span"""
        attrs = {"iteration": iteration, **attributes}
        with self.span("reasoning.iteration", SpanType.REASONING, **attrs) as s:
            yield s

    @contextmanager
    def task_span(self, session_id: str = "", **attributes: Any) -> Generator[Span, None, None]:
        """完整任务 Span"""
        attrs = {"session_id": session_id, **attributes}
        with self.span("agent.task", SpanType.TASK, **attrs) as s:
            yield s

    # ==================== 非上下文管理器 API ====================
    # 用于 run() 等多返回路径的场景

    def begin_trace(self, session_id: str, metadata: dict[str, Any] | None = None) -> Trace | None:
        """
        开始一个新的 Trace（非上下文管理器版本）。

        必须手动调用 end_trace() 来结束。
        适合多 return 路径的长方法。
        """
        if not self._enabled:
            return None

        trace = Trace(
            trace_id=str(uuid.uuid4()),
            session_id=session_id,
            start_time=time.time(),
            metadata=metadata or {},
        )
        self._current_trace = trace
        self._span_stack = []
        return trace

    def end_trace(self, metadata: dict[str, Any] | None = None) -> None:
        """
        结束当前 Trace（非上下文管理器版本）。

        与 begin_trace() 配对使用。
        """
        if not self._enabled or not self._current_trace:
            return

        if metadata:
            self._current_trace.metadata.update(metadata)

        self._current_trace.finish()
        self._export_trace(self._current_trace)
        self._current_trace = None
        self._span_stack = []

    def _export_trace(self, trace: Trace) -> None:
        """导出 Trace 到所有已注册的导出器"""
        for exporter in self._exporters:
            try:
                exporter.export(trace)
            except Exception as e:
                logger.warning(f"[Tracing] Failed to export trace to {type(exporter).__name__}: {e}")


# 全局 tracer 实例
_global_tracer: AgentTracer | None = None


def get_tracer() -> AgentTracer:
    """获取全局 tracer 实例"""
    global _global_tracer
    if _global_tracer is None:
        _global_tracer = AgentTracer(enabled=False)  # 默认禁用，需要显式启用
    return _global_tracer


def set_tracer(tracer: AgentTracer) -> None:
    """设置全局 tracer 实例"""
    global _global_tracer
    _global_tracer = tracer
