# -*- coding: utf-8 -*-
"""
TraceContext 与 TraceSpan 数据模型

基于 contextvars 的请求级追踪上下文，存储 request_id 和有序 span 列表。
支持 HTTP / SSE / WebSocket / 后台 Job 四种 trace 类型，
每种类型使用不同的 request_id 前缀以便区分。
"""

from __future__ import annotations

import contextvars
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any


# ──────────────────────────────────────────────
# Span 类型常量（字符串常量类，不用 Enum）
# ──────────────────────────────────────────────

class SpanType:
    """所有支持的 span 类型常量。"""

    HTTP_IN = "HTTP_IN"
    AUTH = "AUTH"
    ROUTE = "ROUTE"
    SERVICE = "SERVICE"
    DB_QUERY = "DB_QUERY"
    DB_CONN = "DB_CONN"
    DB_CONN_RELEASE = "DB_CONN_RELEASE"
    HTTP_OUT = "HTTP_OUT"
    ERROR = "ERROR"
    DB_ERROR = "DB_ERROR"
    MIDDLEWARE = "MIDDLEWARE"
    MIDDLEWARE_ERROR = "MIDDLEWARE_ERROR"
    SSE_START = "SSE_START"
    SSE_EVENT = "SSE_EVENT"
    SSE_END = "SSE_END"
    AI_CALL = "AI_CALL"
    AI_STREAM = "AI_STREAM"
    AI_ERROR = "AI_ERROR"
    WS_CONNECT = "WS_CONNECT"
    WS_MESSAGE = "WS_MESSAGE"
    WS_DISCONNECT = "WS_DISCONNECT"
    JOB_START = "JOB_START"
    JOB_END = "JOB_END"
    JOB_ERROR = "JOB_ERROR"


# ──────────────────────────────────────────────
# Trace 类型常量
# ──────────────────────────────────────────────

class TraceType:
    """Trace 类型：HTTP / SSE / WebSocket / 后台 Job。"""

    HTTP = "http"
    SSE = "sse"
    WS = "ws"
    JOB = "job"


# ──────────────────────────────────────────────
# 数据模型
# ──────────────────────────────────────────────

@dataclass
class TraceSpan:
    """单个追踪节点，记录某一层的函数调用信息。"""

    span_type: str          # SpanType 中的常量值
    module: str             # 模块路径，如 "xcx_tasks"
    function: str           # 函数名，如 "get_task_list"
    description_zh: str     # 中文描述
    description_en: str     # 英文描述
    params: dict[str, Any]  # 调用参数
    result_summary: str     # 结果摘要
    duration_ms: float      # 耗时（毫秒）
    timestamp: str          # ISO 格式时间戳
    extra: dict[str, Any] = field(default_factory=dict)  # 额外信息（SQL 等）


@dataclass
class TraceContext:
    """请求级追踪上下文，维护 request_id 和有序 span 列表。"""

    request_id: str = field(default_factory=lambda: uuid.uuid4().hex[:12])
    trace_type: str = TraceType.HTTP  # http / sse / ws / job
    start_time: datetime = field(default_factory=datetime.now)
    method: str = ""
    path: str = ""
    user_id: int | None = None
    site_id: int | None = None
    spans: list[TraceSpan] = field(default_factory=list)

    def add_span(self, span: TraceSpan) -> None:
        """追加 span 到列表，保持插入顺序。"""
        self.spans.append(span)


# ──────────────────────────────────────────────
# contextvars 存储
# ──────────────────────────────────────────────

trace_context_var: contextvars.ContextVar[TraceContext | None] = contextvars.ContextVar(
    "trace_context", default=None
)


def get_current_trace() -> TraceContext | None:
    """获取当前请求的 TraceContext（无则返回 None）。"""
    return trace_context_var.get()


def set_current_trace(ctx: TraceContext) -> contextvars.Token:
    """设置当前请求的 TraceContext，返回 Token 用于后续恢复。"""
    return trace_context_var.set(ctx)


# ──────────────────────────────────────────────
# 工厂函数：创建不同类型的 TraceContext
# ──────────────────────────────────────────────

def create_http_trace(method: str, path: str) -> TraceContext:
    """创建 HTTP 请求的 TraceContext（request_id = uuid hex[:12]）。"""
    return TraceContext(
        request_id=uuid.uuid4().hex[:12],
        trace_type=TraceType.HTTP,
        method=method,
        path=path,
    )


def create_ws_trace() -> TraceContext:
    """创建 WebSocket 连接的 TraceContext（request_id = ws_ 前缀）。"""
    return TraceContext(
        request_id=f"ws_{uuid.uuid4().hex[:12]}",
        trace_type=TraceType.WS,
    )


def create_job_trace(job_name: str) -> TraceContext:
    """创建后台 Job 的 TraceContext（request_id = job_ 前缀）。"""
    return TraceContext(
        request_id=f"job_{uuid.uuid4().hex[:12]}",
        trace_type=TraceType.JOB,
        path=job_name,
    )