"""
tracer.py
=========
专家预取可观测性基础设施 — 结构化 JSONL 事件流写入器

Issue #1: 专家预取追踪探针

每行一个事件对象，字段说明：
  ts          : Unix 时间戳（秒，浮点）
  event       : 事件类型（见下文）
  expert_id   : 专家 ID
  layer       : 层索引（-1 = 未知）
  context     : "hot" | "cold" | "gate"
  step        : decode step 编号
  phase       : "draft" | "verify" | "ar"
  draft_pos   : 草稿深度（draft 阶段使用，None 时不写入）

支持的事件类型：
  prefetch_start   — 专家 H2D 传输开始
  prefetch_ready   — 专家 H2D 传输完成（含 transfer_ms）
  expert_evict     — 专家从 GPU 回收（含 was_used）
  expert_wait_hit  — wait_expert() 时专家已 READY
  expert_wait_miss — wait_expert() 时需同步等待
  ct_update        — CorrelationTable.update() 调用
  ct_predict       — CorrelationTable.predict_hot_experts() 调用
"""

from __future__ import annotations

import json
import threading
import time
from typing import Any, Dict, Optional


class ExpertTracer:
    """
    线程安全的 JSONL 专家追踪事件写入器。

    Args:
        output_path : 输出文件路径（JSONL 格式，每行一个 JSON 对象）
        enabled     : 是否启用（False 时 emit() 调用零开销，不加锁不写入）
    """

    def __init__(self, output_path: str, enabled: bool = True):
        self._enabled = enabled

        # 当前推理阶段（由外部调用者通过 set_phase() 更新）
        self._phase: str = "ar"
        self._step: int = 0
        self._draft_pos: Optional[int] = None

        if enabled:
            self._file = open(output_path, "w", buffering=8192)  # noqa: WPS515
            self._lock = threading.Lock()
        else:
            self._file = None  # type: ignore[assignment]
            self._lock = None  # type: ignore[assignment]

    # ------------------------------------------------------------------
    # Phase 设置接口
    # ------------------------------------------------------------------

    def set_phase(
        self,
        phase: str,
        step: int,
        draft_pos: Optional[int] = None,
    ) -> None:
        """
        设置当前推理阶段（由 EAGLE3SDDriver 或主推理循环调用）。

        Args:
            phase     : "draft" | "verify" | "ar"
            step      : 当前 decode step 编号（0-indexed）
            draft_pos : 草稿深度（仅 draft 阶段使用，其余设为 None）
        """
        self._phase = phase
        self._step = step
        self._draft_pos = draft_pos

    # ------------------------------------------------------------------
    # 写入接口
    # ------------------------------------------------------------------

    def emit(self, event_dict: Dict[str, Any]) -> None:
        """
        线程安全地写入一条 JSONL 事件。

        enabled=False 时此方法**立即返回**，无任何锁操作（零开销）。
        """
        if not self._enabled:
            return
        event_dict["ts"] = time.time()
        line = json.dumps(event_dict, ensure_ascii=False) + "\n"
        with self._lock:
            self._file.write(line)

    # ------------------------------------------------------------------
    # 生命周期
    # ------------------------------------------------------------------

    def flush(self) -> None:
        """刷新写缓冲区（不关闭文件）。"""
        if self._enabled and self._file:
            with self._lock:
                self._file.flush()

    def close(self) -> None:
        """刷新缓冲区并关闭文件句柄。"""
        if self._enabled and self._file:
            with self._lock:
                self._file.flush()
                self._file.close()
                self._file = None  # type: ignore[assignment]

    def __del__(self) -> None:
        try:
            self.close()
        except Exception:
            pass

    # ------------------------------------------------------------------
    # 便捷属性（供 CorrelationTable 等读取当前上下文）
    # ------------------------------------------------------------------

    @property
    def current_phase(self) -> str:
        return self._phase

    @property
    def current_step(self) -> int:
        return self._step

    @property
    def current_draft_pos(self) -> Optional[int]:
        return self._draft_pos
