"""
Qihu-AIFlow 统一可观测基座 (Phase 2.1)
=====================================
提供贯穿一次请求的全局链路 ID（trace_id）与结构化日志能力，
为后续重构（上下文引擎、媒体/LLM 成本归因）建立回归基线。

设计原则（Strangler Fig 增量）：
- 纯增量埋点，不改任何既有接口的返回结构。
- 向后兼容现有 `logging.getLogger(__name__)` 用法：trace_id 由 root logger 的
  TraceIdFilter 自动注入，调用方无需改动。
- 不引入第三方依赖（structlog 未安装），自带可选的 JSON formatter。

使用方式：
    from app.core.observability import set_trace_id, get_trace_id, span

    # 在中间件里为每个请求注入：
    set_trace_id(request_id)

    # 在关键调用处打计时 span（自动带 trace_id）：
    async with span("chat.send_message", conv_id=conv_id):
        return await _send_message_core(...)
"""
import logging
import time
import uuid
import json
from contextvars import ContextVar

# ---------------------------------------------------------------------------
# trace_id：贯穿一次请求（含其内所有子调用/协程）的全局链路 ID
# ---------------------------------------------------------------------------
_trace_id_var: ContextVar[str] = ContextVar("aiflow_trace_id", default="-")


def new_trace_id() -> str:
    """生成一个崭新的链路 ID（32 位 hex）。"""
    return uuid.uuid4().hex


def set_trace_id(tid: str | None = None) -> str:
    """设置当前上下文的 trace_id；不传则自动生成。返回最终使用的 ID。"""
    tid = tid or new_trace_id()
    _trace_id_var.set(tid)
    return tid


def get_trace_id() -> str:
    """读取当前上下文的 trace_id；未设置时返回占位 '-'。"""
    return _trace_id_var.get()


class TraceIdFilter(logging.Filter):
    """为每一条日志记录注入 trace_id 属性（不存在则写为 '-'）。

    挂载到 root logger 后，所有 `getLogger(__name__)` 的子 logger 都会继承，
    从而让全平台既有日志自动带上链路 ID，零侵入。
    """

    def filter(self, record: logging.LogRecord) -> bool:
        if not hasattr(record, "trace_id"):
            try:
                record.trace_id = get_trace_id()
            except Exception:
                record.trace_id = "-"
        return True


# 标准 logging 字段，JSON 序列化时需排除，避免与自定义 extra 冲突
_STD_ATTRS = set(logging.makeLogRecord({}).__dict__.keys()) | {
    "message", "asctime", "msg", "trace_id", "exc_info", "stack_info", "args",
}


class JsonFormatter(logging.Formatter):
    """结构化 JSON 日志格式（通过环境变量 AIFLOW_JSON_LOGS=1 启用）。

    默认不启用，避免破坏现有日志采集/人工查阅习惯。
    """

    def format(self, record: logging.LogRecord) -> str:
        data = {
            "ts": self.formatTime(record, "%Y-%m-%dT%H:%M:%S%z"),
            "level": record.levelname,
            "logger": record.name,
            "trace_id": getattr(record, "trace_id", "-"),
            "msg": record.getMessage(),
        }
        for k, v in record.__dict__.items():
            if k not in _STD_ATTRS:
                try:
                    json.dumps(v, ensure_ascii=False)
                    data[k] = v
                except (TypeError, ValueError):
                    data[k] = str(v)
        if record.exc_info:
            data["exc"] = self.formatException(record.exc_info)
        return json.dumps(data, ensure_ascii=False)


def get_logger(name: str | None = None) -> logging.Logger:
    """获取 logger（向后兼容现有用法）。trace_id 由 root filter 自动注入。"""
    return logging.getLogger(name)


# ---------------------------------------------------------------------------
# span：结构化计时上下文管理器（耗时 + trace_id + 自定义字段）
# ---------------------------------------------------------------------------
class _Span:
    """同步/异步通用的计时 span。退出时记录 elapsed_ms 与状态。"""

    def __init__(self, name: str, **fields):
        self.name = name
        self.fields = fields
        self._logger = logging.getLogger("aiflow.span")
        self._start = None
        self.elapsed_ms = 0.0

    def _enter(self):
        self._start = time.perf_counter()
        self._logger.info(
            f"SPAN_START {self.name}",
            extra={"span": self.name, "event": "start", "trace_id": get_trace_id(), **self.fields},
        )
        return self

    def _exit(self, exc_type, exc, tb):
        self.elapsed_ms = (time.perf_counter() - self._start) * 1000
        status = "error" if exc_type else "ok"
        self._logger.info(
            f"SPAN_END {self.name} {self.elapsed_ms:.1f}ms",
            extra={
                "span": self.name,
                "event": "end",
                "elapsed_ms": round(self.elapsed_ms, 1),
                "status": status,
                "trace_id": get_trace_id(),
                **self.fields,
            },
        )
        return False  # 不吞掉异常

    # 同步协议
    def __enter__(self):
        return self._enter()

    def __exit__(self, exc_type, exc, tb):
        return self._exit(exc_type, exc, tb)

    # 异步协议
    async def __aenter__(self):
        return self._enter()

    async def __aexit__(self, exc_type, exc, tb):
        return self._exit(exc_type, exc, tb)


def span(name: str, **fields):
    """创建一个计时 span。可同步或异步 with。

    async with span("chat.send_message", conv_id=cid):
        ...
    """
    return _Span(name, **fields)


__all__ = [
    "new_trace_id",
    "set_trace_id",
    "get_trace_id",
    "TraceIdFilter",
    "JsonFormatter",
    "get_logger",
    "span",
]