#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""作业14 · 从 Langfuse 导出 JSON/CSV 汇总「昨日 Top10 烧 Token 的应用」

支持输入：
  1) Langfuse UI 导出的 traces JSON（数组或 {data:[...]}）
  2) CSV（表头可含本作业字段规范名，或常见别名）

用法：
  python hw14_02_昨日Top10烧Token.py --input data/homework14/sample_export.json
  python hw14_02_昨日Top10烧Token.py --input export.csv --date 2026-08-26
  python hw14_02_昨日Top10烧Token.py --input export.json --tz Asia/Shanghai

输出：
  data/homework14/top10_tokens_YYYYMMDD.csv
  data/homework14/top10_tokens_YYYYMMDD.md
  终端打印 Top10 表
"""
from __future__ import annotations

import argparse
import csv
import json
import os
from collections import defaultdict
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any


FIELD_SPEC = [
    "trace_id",
    "session_id",
    "model",
    "prompt_tokens",
    "completion_tokens",
    "tool_name",
    "latency_ms",
    "error_type",
    "app_name",  # 汇总维度（规范扩展：应用名）
    "timestamp",
]

# 导出文件里可能出现的别名 → 规范字段
ALIASES = {
    "trace_id": ["trace_id", "traceId", "id", "TraceId"],
    "session_id": ["session_id", "sessionId", "session_id", "userId", "user_id", "session"],
    "model": ["model", "model_name", "modelName", "Model"],
    "prompt_tokens": [
        "prompt_tokens",
        "promptTokens",
        "input_tokens",
        "inputTokens",
        "message_tokens",
    ],
    "completion_tokens": [
        "completion_tokens",
        "completionTokens",
        "output_tokens",
        "outputTokens",
        "answer_tokens",
    ],
    "tool_name": ["tool_name", "toolName", "name", "span_name", "observation_name"],
    "latency_ms": ["latency_ms", "latencyMs", "latency", "duration_ms", "durationMs", "duration"],
    "error_type": ["error_type", "errorType", "level", "statusMessage", "error"],
    "app_name": [
        "app_name",
        "appName",
        "application",
        "tags",
        "metadata.app",
        "metadata.app_name",
        "name",
    ],
    "timestamp": [
        "timestamp",
        "createdAt",
        "created_at",
        "startTime",
        "start_time",
        "time",
    ],
}


def out_dir() -> Path:
    env = os.environ.get("HW14_OUT")
    if env:
        p = Path(env)
        p.mkdir(parents=True, exist_ok=True)
        return p
    repo = Path(__file__).resolve().parents[2] / "data" / "homework14"
    candidates = [repo]
    opt = Path("/opt/workbuddy/data/homework14")
    if os.name != "nt" and Path("/opt/workbuddy").is_dir():
        candidates.insert(0, opt)
    for c in candidates:
        try:
            c.mkdir(parents=True, exist_ok=True)
            return c
        except OSError:
            continue
    p = Path.cwd() / "data" / "homework14"
    p.mkdir(parents=True, exist_ok=True)
    return p


def dig(obj: Any, dotted: str) -> Any:
    cur = obj
    for part in dotted.split("."):
        if not isinstance(cur, dict) or part not in cur:
            return None
        cur = cur[part]
    return cur


def pick(row: dict, canonical: str) -> Any:
    for key in ALIASES.get(canonical, [canonical]):
        if "." in key:
            v = dig(row, key)
        else:
            v = row.get(key)
        if v is None or v == "":
            continue
        if canonical == "app_name" and isinstance(v, list):
            # tags: ["app:客服助手", ...]
            for t in v:
                s = str(t)
                if s.startswith("app:"):
                    return s.split(":", 1)[1]
            return str(v[0]) if v else "unknown"
        return v
    # metadata nested
    meta = row.get("metadata") if isinstance(row.get("metadata"), dict) else {}
    if canonical == "app_name":
        for k in ("app", "app_name", "application", "appName"):
            if meta.get(k):
                return meta[k]
    if canonical == "model" and meta.get("model"):
        return meta["model"]
    return None


def to_int(v: Any) -> int:
    if v is None or v == "":
        return 0
    try:
        return int(float(v))
    except Exception:
        return 0


def to_float(v: Any) -> float | None:
    if v is None or v == "":
        return None
    try:
        x = float(v)
        # 若像秒（很小）则转 ms：启发式 < 100 且非整数大 latency
        return x
    except Exception:
        return None


def parse_ts(v: Any) -> datetime | None:
    if v is None or v == "":
        return None
    if isinstance(v, (int, float)):
        # ms or s
        x = float(v)
        if x > 1e12:
            x /= 1000.0
        elif x > 1e10:
            x /= 1000.0
        return datetime.fromtimestamp(x)
    s = str(v).replace("Z", "+00:00")
    for fmt in (
        None,  # fromisoformat
        "%Y-%m-%d %H:%M:%S",
        "%Y-%m-%dT%H:%M:%S",
        "%Y/%m/%d %H:%M:%S",
    ):
        try:
            if fmt is None:
                return datetime.fromisoformat(s)
            return datetime.strptime(s[:19], fmt)
        except Exception:
            continue
    return None


def load_rows(path: Path) -> list[dict]:
    text = path.read_text(encoding="utf-8-sig")
    if path.suffix.lower() == ".csv":
        r = csv.DictReader(text.splitlines())
        return [dict(x) for x in r]
    data = json.loads(text)
    if isinstance(data, list):
        return data
    if isinstance(data, dict):
        for k in ("data", "traces", "items", "rows"):
            if isinstance(data.get(k), list):
                return data[k]
        # single object
        return [data]
    raise ValueError(f"无法解析: {path}")


def normalize(row: dict) -> dict:
    prompt = to_int(pick(row, "prompt_tokens"))
    completion = to_int(pick(row, "completion_tokens"))
    # Langfuse 有时只有 total
    total = to_int(row.get("totalTokens") or row.get("total_tokens") or 0)
    if prompt == 0 and completion == 0 and total:
        prompt = total
    latency = to_float(pick(row, "latency_ms"))
    if latency is not None and latency < 50 and latency == int(latency) is False:
        # 可能是秒
        if latency < 100:
            latency = latency * 1000.0
    app = pick(row, "app_name") or pick(row, "tool_name") or "unknown"
    # 若 name 是 message 一类，尽量用 metadata
    if str(app).lower() in ("message", "trace", "generation", "span", "unknown"):
        meta = row.get("metadata") if isinstance(row.get("metadata"), dict) else {}
        app = meta.get("app_name") or meta.get("app") or meta.get("application") or app

    return {
        "trace_id": str(pick(row, "trace_id") or ""),
        "session_id": str(pick(row, "session_id") or ""),
        "model": str(pick(row, "model") or "unknown"),
        "prompt_tokens": prompt,
        "completion_tokens": completion,
        "total_tokens": prompt + completion if (prompt or completion) else total,
        "tool_name": str(pick(row, "tool_name") or ""),
        "latency_ms": latency,
        "error_type": str(pick(row, "error_type") or "") if pick(row, "error_type") else "",
        "app_name": str(app),
        "timestamp": parse_ts(pick(row, "timestamp")),
    }


def filter_day(rows: list[dict], day: datetime) -> list[dict]:
    start = day.replace(hour=0, minute=0, second=0, microsecond=0)
    end = start + timedelta(days=1)
    out = []
    for r in rows:
        ts = r["timestamp"]
        if ts is None:
            # 无时间戳：若用户指定 --include-missing-ts 才纳入；默认跳过
            continue
        # 去掉 tzinfo 做日期比较
        if ts.tzinfo is not None:
            ts = ts.replace(tzinfo=None)
        if start <= ts < end:
            out.append(r)
    return out


def top10(rows: list[dict]) -> list[dict]:
    agg: dict[str, dict] = {}
    for r in rows:
        app = r["app_name"] or "unknown"
        a = agg.setdefault(
            app,
            {
                "app_name": app,
                "traces": 0,
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0,
                "errors": 0,
                "latency_sum": 0.0,
                "latency_n": 0,
            },
        )
        a["traces"] += 1
        a["prompt_tokens"] += r["prompt_tokens"]
        a["completion_tokens"] += r["completion_tokens"]
        a["total_tokens"] += r["total_tokens"]
        if r["error_type"]:
            a["errors"] += 1
        if r["latency_ms"] is not None:
            a["latency_sum"] += float(r["latency_ms"])
            a["latency_n"] += 1
    items = list(agg.values())
    for a in items:
        a["latency_avg_ms"] = (
            round(a["latency_sum"] / a["latency_n"], 1) if a["latency_n"] else None
        )
        del a["latency_sum"]
        del a["latency_n"]
    items.sort(key=lambda x: x["total_tokens"], reverse=True)
    return items[:10]


def main() -> int:
    ap = argparse.ArgumentParser(description="作业14 昨日 Top10 Token")
    ap.add_argument("--input", required=True, help="导出 JSON 或 CSV 路径")
    ap.add_argument("--date", default="", help="YYYY-MM-DD，默认昨天（本地）")
    ap.add_argument("--top", type=int, default=10)
    ap.add_argument(
        "--include-missing-ts",
        action="store_true",
        help="无 timestamp 的行也计入（样本练习可用）",
    )
    args = ap.parse_args()

    path = Path(args.input)
    if not path.is_file():
        raise SystemExit(f"[FAIL] 文件不存在: {path}")

    if args.date:
        day = datetime.strptime(args.date, "%Y-%m-%d")
    else:
        day = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0) - timedelta(days=1)

    raw = load_rows(path)
    norm = [normalize(r) for r in raw]
    if args.include_missing_ts:
        for r in norm:
            if r["timestamp"] is None:
                r["timestamp"] = day.replace(hour=12)
    day_rows = filter_day(norm, day)
    ranked = top10(day_rows)[: args.top]

    day_s = day.strftime("%Y-%m-%d")
    print(f"[INFO] 输入={path} 行数={len(raw)} 规范化={len(norm)} 昨日({day_s})命中={len(day_rows)}")
    print(f"\n=== 昨日 Top{args.top} 烧 Token 的应用 ===")
    print(f"{'rank':<4} {'app_name':<28} {'total_tokens':>12} {'traces':>8} {'errors':>7}")
    for i, a in enumerate(ranked, 1):
        print(
            f"{i:<4} {a['app_name'][:28]:<28} {a['total_tokens']:>12} {a['traces']:>8} {a['errors']:>7}"
        )
    if not ranked:
        print("[WARN] 无数据。可检查日期，或加 --include-missing-ts，或使用 sample_export.json")

    out = out_dir()
    stamp = day.strftime("%Y%m%d")
    csv_path = out / f"top10_tokens_{stamp}.csv"
    md_path = out / f"top10_tokens_{stamp}.md"
    with csv_path.open("w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(
            f,
            fieldnames=[
                "rank",
                "app_name",
                "total_tokens",
                "prompt_tokens",
                "completion_tokens",
                "traces",
                "errors",
                "latency_avg_ms",
            ],
        )
        w.writeheader()
        for i, a in enumerate(ranked, 1):
            w.writerow({"rank": i, **a})

    lines = [
        f"# 昨日 Top{args.top} 烧 Token 的应用（{day_s}）",
        "",
        f"数据源：`{path}`",
        "",
        "| rank | app_name | total_tokens | traces | errors | latency_avg_ms |",
        "|---:|---|---:|---:|---:|---:|",
    ]
    for i, a in enumerate(ranked, 1):
        lines.append(
            f"| {i} | {a['app_name']} | {a['total_tokens']} | {a['traces']} | {a['errors']} | {a['latency_avg_ms']} |"
        )
    lines.append("")
    lines.append("字段规范：" + ", ".join(FIELD_SPEC))
    md_path.write_text("\n".join(lines), encoding="utf-8")
    print(f"\n[OK] {csv_path}")
    print(f"[OK] {md_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())