#!/usr/bin/env python3
from __future__ import annotations

import argparse
import csv
import json
from collections import Counter
from pathlib import Path
from typing import Any


def main() -> int:
    parser = argparse.ArgumentParser(description="Summarize a CML eval workflow metric_diff directory.")
    parser.add_argument("--workflow-root", default="/mnt/xiaoai-zk-model-train-tj5/workflow5")
    parser.add_argument("--run-dic", required=True, help="workflow runDic, for example 18065")
    parser.add_argument("--output", help="Markdown output path. Prints to stdout when omitted.")
    parser.add_argument("--max-error-examples", type=int, default=30)
    args = parser.parse_args()

    workflow_dir = Path(args.workflow_root) / f"workflow{args.run_dic}"
    metric_dir = workflow_dir / "metric_diff"
    result = summarize(metric_dir, args.max_error_examples)
    text = render_markdown(args.run_dic, workflow_dir, result)
    if args.output:
        out = Path(args.output)
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(text, encoding="utf-8")
    else:
        print(text)
    return 0 if result["ok"] else 1


def summarize(metric_dir: Path, max_error_examples: int) -> dict[str, Any]:
    lark_path = metric_dir / "lark_template.json"
    comparison_path = metric_dir / "specific_comparison.csv"
    result: dict[str, Any] = {
        "ok": True,
        "metric_dir": str(metric_dir),
        "missing": [],
        "metrics": [],
        "specific_rows": 0,
        "error_summary": [],
        "error_examples": [],
    }
    if not lark_path.exists():
        result["missing"].append(str(lark_path))
    else:
        result["metrics"] = extract_lark_metrics(load_json(lark_path))
    if comparison_path.exists():
        rows = read_csv(comparison_path)
        result["specific_rows"] = len(rows)
        result["error_summary"] = summarize_comparison(rows)
        result["error_examples"] = select_error_examples(rows, max_error_examples)
    else:
        result["missing"].append(str(comparison_path))
    result["ok"] = not result["missing"]
    return result


def load_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


def extract_lark_metrics(obj: Any) -> list[dict[str, Any]]:
    metrics: list[dict[str, Any]] = []

    def walk(value: Any, path: list[str]) -> None:
        if isinstance(value, dict):
            maybe_name = value.get("name") or value.get("title") or value.get("sub_cate") or value.get("cate")
            maybe_acc = first_present(value, ["acc", "accuracy", "new_acc", "dev_acc", "rate", "pass_rate"])
            maybe_total = first_present(value, ["total", "count", "all"])
            maybe_right = first_present(value, ["right", "correct", "hit"])
            if maybe_acc is not None or (maybe_total is not None and maybe_right is not None):
                metrics.append(
                    {
                        "path": " / ".join(path + ([str(maybe_name)] if maybe_name else [])),
                        "acc": maybe_acc,
                        "right": maybe_right,
                        "total": maybe_total,
                    }
                )
            for key, child in value.items():
                if isinstance(child, (dict, list)):
                    walk(child, path + [str(key)])
        elif isinstance(value, list):
            for idx, child in enumerate(value):
                if isinstance(child, (dict, list)):
                    walk(child, path + [str(idx)])

    walk(obj, [])
    # De-duplicate while preserving order.
    seen: set[tuple[str, str, str, str]] = set()
    uniq: list[dict[str, Any]] = []
    for item in metrics:
        key = tuple(str(item.get(k, "")) for k in ("path", "acc", "right", "total"))
        if key not in seen:
            seen.add(key)
            uniq.append(item)
    return uniq[:80]


def first_present(obj: dict[str, Any], keys: list[str]) -> Any:
    for key in keys:
        if key in obj and obj[key] not in ("", None):
            return obj[key]
    return None


def read_csv(path: Path) -> list[dict[str, str]]:
    for encoding in ("utf-8-sig", "utf-8", "gb18030"):
        try:
            with path.open(encoding=encoding, newline="") as handle:
                return list(csv.DictReader(handle))
        except UnicodeDecodeError:
            continue
    raise UnicodeDecodeError("csv", b"", 0, 1, f"cannot decode {path}")


def summarize_comparison(rows: list[dict[str, str]]) -> list[dict[str, Any]]:
    if not rows:
        return []
    fieldnames = set(rows[0])
    cate_field = pick_field(fieldnames, ["sub_cate", "子集", "cate", "类别"])
    gold_field = pick_field(fieldnames, ["code_label", "gold", "label", "类别真实标签", "code标签"])
    pred_field = pick_field(fieldnames, ["origin_predict_dev", "predict_dev", "prediction", "pred", "模型输出"])
    if not cate_field:
        cate_field = "__all__"
    counter: Counter[tuple[str, str, str]] = Counter()
    for row in rows:
        cate = row.get(cate_field, "ALL") if cate_field != "__all__" else "ALL"
        gold = row.get(gold_field, "") if gold_field else ""
        pred = row.get(pred_field, "") if pred_field else ""
        is_error = detect_error(row)
        if is_error:
            counter[(cate, gold[:80], pred[:80])] += 1
    return [
        {"sub_cate": cate, "gold": gold, "pred": pred, "count": count}
        for (cate, gold, pred), count in counter.most_common(30)
    ]


def detect_error(row: dict[str, str]) -> bool:
    for key in ("is_correct", "correct", "是否正确", "same", "is_same"):
        value = row.get(key)
        if value is None:
            continue
        text = str(value).strip().lower()
        if text in {"false", "0", "错", "错误", "no", "n"}:
            return True
        if text in {"true", "1", "对", "正确", "yes", "y"}:
            return False
    base = row.get("origin_predict_base") or row.get("predict_base") or ""
    dev = row.get("origin_predict_dev") or row.get("predict_dev") or row.get("prediction") or ""
    gold = row.get("code_label") or row.get("gold") or row.get("code标签") or ""
    return bool(gold and dev and normalize_label(gold) != normalize_label(dev)) or bool(base and dev and base != dev)


def normalize_label(text: str) -> str:
    return "".join(str(text).split()).replace("'", '"')


def select_error_examples(rows: list[dict[str, str]], limit: int) -> list[dict[str, str]]:
    examples: list[dict[str, str]] = []
    for row in rows:
        if not detect_error(row):
            continue
        examples.append(
            {
                "query": row.get("query") or row.get("当前query") or row.get("current_query") or "",
                "sub_cate": row.get("sub_cate") or row.get("子集") or "",
                "gold": row.get("code_label") or row.get("gold") or row.get("code标签") or "",
                "pred": row.get("origin_predict_dev") or row.get("predict_dev") or row.get("prediction") or "",
            }
        )
        if len(examples) >= limit:
            break
    return examples


def pick_field(fieldnames: set[str], candidates: list[str]) -> str:
    for candidate in candidates:
        if candidate in fieldnames:
            return candidate
    return ""


def render_markdown(run_dic: str, workflow_dir: Path, result: dict[str, Any]) -> str:
    lines: list[str] = [
        f"# workflow{run_dic} 评测摘要",
        "",
        f"- workflow_dir: `{workflow_dir}`",
        f"- metric_dir: `{result['metric_dir']}`",
    ]
    if result["missing"]:
        lines.append(f"- missing: `{', '.join(result['missing'])}`")
    lines += ["", "## 指标摘录", ""]
    if result["metrics"]:
        lines.append("| path | acc | right | total |")
        lines.append("|---|---:|---:|---:|")
        for item in result["metrics"]:
            lines.append(
                f"| {safe_cell(item.get('path'))} | {safe_cell(item.get('acc'))} | {safe_cell(item.get('right'))} | {safe_cell(item.get('total'))} |"
            )
    else:
        lines.append("未解析到 lark_template 指标。")
    lines += ["", "## 错误分布 Top", ""]
    if result["error_summary"]:
        lines.append("| sub_cate | gold | pred | count |")
        lines.append("|---|---|---|---:|")
        for item in result["error_summary"]:
            lines.append(
                f"| {safe_cell(item['sub_cate'])} | {safe_cell(item['gold'])} | {safe_cell(item['pred'])} | {item['count']} |"
            )
    else:
        lines.append("未解析到错误分布。")
    lines += ["", "## 错误样例", ""]
    if result["error_examples"]:
        lines.append("| sub_cate | query | gold | pred |")
        lines.append("|---|---|---|---|")
        for item in result["error_examples"]:
            lines.append(
                f"| {safe_cell(item['sub_cate'])} | {safe_cell(item['query'])} | {safe_cell(item['gold'])} | {safe_cell(item['pred'])} |"
            )
    else:
        lines.append("未抽到错误样例。")
    lines.append("")
    return "\n".join(lines)


def safe_cell(value: Any) -> str:
    text = str(value or "").replace("\n", "<br>").replace("|", "\\|")
    return text[:500]


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