"""
自建的轻量级 tracing / 可观测系统。

设计沿用分布式追踪的 span 树模型（见书 6.x「Agent 的可观测性」）：
    - 一次 agent 任务 = 一条 Trace
    - 每次 LLM 调用 / 工具调用 = 一个 Span
    - Span 记录：所属步骤、类型、token 用量（prompt/completion/cached）、时延、成本

用法：
    tracer = Tracer(client)
    resp = tracer.chat(step="turn-1", tool="query_order",
                       model=..., messages=..., temperature=0)
    tracer.print_breakdown()   # 打印按步骤/工具聚合的成本拆解

离线复用（不打模型、只算成本）：
    tracer = Tracer.from_records(records, pricing=..., name=...)
    # records 里是此前真实运行录下的每一步 token 用量（canned token counts）
"""

import math
import time
from dataclasses import dataclass, asdict, field
from typing import List, Optional

from config import Pricing, default_pricing


def _percentile(values: List[float], q: float) -> float:
    """最近秩（nearest-rank）百分位，避免引入 numpy 依赖。q 取 0~100。"""
    if not values:
        return 0.0
    xs = sorted(values)
    if len(xs) == 1:
        return xs[0]
    # 最近秩 = ceil(q/100 * N)。不能用 int(round(x + 0.5))：当 q/100*N 恰为
    # 整数 k 时，round(k + 0.5) 的银行家舍入会得到 k+1（如 n=100、q=99 时
    # 秩变成 100，把 p99 报成最大值）。
    rank = max(1, min(len(xs), math.ceil(q / 100.0 * len(xs))))
    return xs[rank - 1]


@dataclass
class Span:
    """一次被追踪的调用（这里主要是 LLM 调用）。"""
    step: str                 # 逻辑步骤名，如 "turn-2"
    tool: str                 # 该步骤关联的工具/动作名，用于归因“哪一步最贵”
    kind: str = "llm"         # span 类型：llm / tool
    prompt_tokens: int = 0
    cached_tokens: int = 0
    completion_tokens: int = 0
    reasoning_tokens: int = 0
    # 该轮输入里「工具返回结果」占用的累计 token（同一份工具返回会在后续每轮被反复计费）。
    # 由上层用 tokenizer 估算并填入；离线复用时从 records 读回。-1 表示未知。
    tool_ctx_tokens: int = -1
    latency_s: float = 0.0
    cost_usd: float = 0.0
    response_id: str = ""
    response_model: str = ""
    response_created: int = 0

    @property
    def total_tokens(self) -> int:
        return self.prompt_tokens + self.completion_tokens

    @property
    def uncached_prompt_tokens(self) -> int:
        return max(self.prompt_tokens - self.cached_tokens, 0)


class Tracer:
    """包裹 OpenAI client，自动记录每次 LLM 调用的 usage / 时延 / 成本。"""

    def __init__(self, client=None, name: str = "trace",
                 pricing: Optional[Pricing] = None):
        self.client = client
        self.name = name
        self.pricing = pricing or default_pricing()
        self.spans: List[Span] = []

    # ---------- 采集 ----------
    def chat(self, step: str, tool: str, tool_ctx_tokens: int = -1, **kwargs):
        """发起一次被追踪的 chat.completions 调用。

        kwargs 原样透传给 openai client（model / messages / temperature 等）。
        tool_ctx_tokens：本轮输入里工具返回结果占用的累计 token（可选，用于成本归因）。
        返回原始的 OpenAI response 对象，方便上层取 content。
        """
        t0 = time.time()
        resp = self.client.chat.completions.create(**kwargs)
        latency = time.time() - t0

        usage = resp.usage
        # Some OpenAI-compatible providers omit usage (null); match from_records coercion.
        if usage is None:
            span = Span(
                step=step,
                tool=tool,
                kind="llm",
                tool_ctx_tokens=tool_ctx_tokens,
                latency_s=latency,
                cost_usd=0.0,
            )
            self.spans.append(span)
            return resp

        # cached_tokens 藏在 prompt_tokens_details 里，注意做防御式读取
        cached = 0
        details = getattr(usage, "prompt_tokens_details", None)
        if details is not None:
            cached = getattr(details, "cached_tokens", 0) or 0

        prompt_tokens = int(getattr(usage, "prompt_tokens", 0) or 0)
        completion_tokens = int(getattr(usage, "completion_tokens", 0) or 0)
        completion_details = getattr(usage, "completion_tokens_details", None)
        reasoning_tokens = int(getattr(completion_details, "reasoning_tokens", 0) or 0)
        span = Span(
            step=step,
            tool=tool,
            kind="llm",
            prompt_tokens=prompt_tokens,
            cached_tokens=cached,
            completion_tokens=completion_tokens,
            reasoning_tokens=reasoning_tokens,
            tool_ctx_tokens=tool_ctx_tokens,
            latency_s=latency,
            cost_usd=self.pricing.cost_usd(
                prompt_tokens, cached, completion_tokens),
            response_id=str(getattr(resp, "id", "") or ""),
            response_model=str(getattr(resp, "model", "") or ""),
            response_created=int(getattr(resp, "created", 0) or 0),
        )
        self.spans.append(span)
        return resp

    # ---------- 离线复用（canned token counts → 重新计成本）----------
    @classmethod
    def from_records(cls, records: List[dict], name: str = "trace",
                     pricing: Optional[Pricing] = None) -> "Tracer":
        """用此前录下的 token 用量重建一条 trace，并按给定单价重算成本（不打模型）。"""
        tr = cls(client=None, name=name, pricing=pricing)
        for r in records:
            span = Span(
                step=r.get("step", ""),
                tool=r.get("tool", ""),
                kind=r.get("kind", "llm"),
                # Missing/null tool_ctx → -1 (unknown); keep explicit 0 (known-zero).
                prompt_tokens=int(r.get("prompt_tokens") or 0),
                cached_tokens=int(r.get("cached_tokens") or 0),
                completion_tokens=int(r.get("completion_tokens") or 0),
                reasoning_tokens=int(r.get("reasoning_tokens") or 0),
                tool_ctx_tokens=(-1 if r.get("tool_ctx_tokens") is None
                                 else int(r.get("tool_ctx_tokens"))),
                latency_s=float(r.get("latency_s") or 0.0),
                response_id=str(r.get("response_id") or ""),
                response_model=str(r.get("response_model") or ""),
                response_created=int(r.get("response_created") or 0),
            )
            span.cost_usd = tr.pricing.cost_usd(
                span.prompt_tokens, span.cached_tokens, span.completion_tokens)
            tr.spans.append(span)
        return tr

    def to_records(self) -> List[dict]:
        """导出每一步的原始 token 用量（用于落盘成 canned trace，供离线复用）。"""
        out = []
        for s in self.spans:
            d = asdict(s)
            d.pop("cost_usd", None)   # 成本由单价重算，不落盘固定值
            out.append(d)
        return out

    # ---------- 聚合 ----------
    def total_cost(self) -> float:
        return sum(s.cost_usd for s in self.spans)

    def total_prompt_tokens(self) -> int:
        return sum(s.prompt_tokens for s in self.spans)

    def total_cached_tokens(self) -> int:
        return sum(s.cached_tokens for s in self.spans)

    def total_completion_tokens(self) -> int:
        return sum(s.completion_tokens for s in self.spans)

    def total_uncached_prompt_tokens(self) -> int:
        return sum(s.uncached_prompt_tokens for s in self.spans)

    def total_tool_ctx_tokens(self) -> int:
        return sum(s.tool_ctx_tokens for s in self.spans if s.tool_ctx_tokens >= 0)

    def total_latency(self) -> float:
        return sum(s.latency_s for s in self.spans)

    def cache_rate(self) -> float:
        pin = self.total_prompt_tokens()
        return self.total_cached_tokens() / pin if pin else 0.0

    def component_costs(self) -> dict:
        """把总成本拆成三个成本构成要素（对应书「成本的构成要素」）：
            - 未缓存输入 / 缓存输入 / 输出
        以及输入侧里「工具返回注入」token 占比（若已知）。"""
        p = self.pricing
        uncached_in = self.total_uncached_prompt_tokens()
        cached_in = self.total_cached_tokens()
        out = self.total_completion_tokens()
        return {
            "uncached_input_cost": uncached_in / 1_000_000 * p.input_per_m,
            "cached_input_cost": cached_in / 1_000_000 * p.cached_per_m,
            "output_cost": out / 1_000_000 * p.output_per_m,
            "uncached_input_tokens": uncached_in,
            "cached_input_tokens": cached_in,
            "output_tokens": out,
            "tool_ctx_tokens": self.total_tool_ctx_tokens(),
        }

    def cost_distribution(self) -> dict:
        """按步骤的单步成本分布（p50/p95/p99）。对应书「成本分布 p50/p95/p99」。"""
        costs = [s.cost_usd for s in self.spans]
        n = len(costs)
        return {
            "n": n,
            "mean": (sum(costs) / n) if n else 0.0,
            "p50": _percentile(costs, 50),
            "p95": _percentile(costs, 95),
            "p99": _percentile(costs, 99),
            "max": max(costs) if costs else 0.0,
        }

    # ---------- 打印 ----------
    def print_breakdown(self, title: Optional[str] = None):
        """打印一次 agent 任务的按步骤成本拆解表，并指出最贵的一步、成本构成与分布。"""
        print()
        print(f"===== 成本拆解: {title or self.name} =====")
        header = (
            f"{'步骤':<8} {'工具/动作':<20} {'输入tok':>8} {'缓存tok':>8} "
            f"{'工具tok':>8} {'输出tok':>8} {'时延(s)':>8} {'成本($)':>12}"
        )
        print(header)
        print("-" * len(header))
        for s in self.spans:
            tctx = s.tool_ctx_tokens if s.tool_ctx_tokens >= 0 else "-"
            print(
                f"{s.step:<8} {s.tool:<20} {s.prompt_tokens:>8} {s.cached_tokens:>8} "
                f"{str(tctx):>8} {s.completion_tokens:>8} {s.latency_s:>8.2f} "
                f"{s.cost_usd:>12.6f}"
            )
        print("-" * len(header))
        tctx_total = self.total_tool_ctx_tokens() if any(
            s.tool_ctx_tokens >= 0 for s in self.spans) else "-"
        print(
            f"{'合计':<8} {'':<20} {self.total_prompt_tokens():>8} "
            f"{self.total_cached_tokens():>8} {str(tctx_total):>8} "
            f"{self.total_completion_tokens():>8} "
            f"{self.total_latency():>8.2f} {self.total_cost():>12.6f}"
        )

        # 归因：哪一步最贵
        if self.spans:
            worst = max(self.spans, key=lambda s: s.cost_usd)
            total = self.total_cost()
            share = worst.cost_usd / total * 100 if total else 0
            print(
                f"\n最贵的一步 → {worst.step} / {worst.tool}: "
                f"${worst.cost_usd:.6f}（占总成本 {share:.1f}%）"
            )

        # 成本构成拆解（未缓存输入 / 缓存输入 / 输出）
        comp = self.component_costs()
        total = self.total_cost() or 1e-12
        print("成本构成:")
        print(f"  未缓存输入  {comp['uncached_input_tokens']:>8} tok  "
              f"${comp['uncached_input_cost']:.6f}  ({comp['uncached_input_cost']/total*100:.1f}%)")
        print(f"  缓存输入    {comp['cached_input_tokens']:>8} tok  "
              f"${comp['cached_input_cost']:.6f}  ({comp['cached_input_cost']/total*100:.1f}%)")
        print(f"  输出        {comp['output_tokens']:>8} tok  "
              f"${comp['output_cost']:.6f}  ({comp['output_cost']/total*100:.1f}%)")
        if comp["tool_ctx_tokens"] > 0:
            print(f"  其中「工具返回注入」累计输入 {comp['tool_ctx_tokens']} tok "
                  f"（同一份工具返回在后续每轮被反复计费）")

        # 单步成本分布
        dist = self.cost_distribution()
        print(f"单步成本分布(n={dist['n']}): 均值 ${dist['mean']:.6f}  "
              f"p50 ${dist['p50']:.6f}  p95 ${dist['p95']:.6f}  p99 ${dist['p99']:.6f}")