from __future__ import annotations

import argparse
import json
import math
from datetime import datetime
from pathlib import Path
from typing import Any

from PIL import Image, ImageDraw, ImageFont


PROJECT_ROOT = Path(__file__).resolve().parents[1]
ANALYSIS_DIR = PROJECT_ROOT / "runs/analysis/validation_error_analysis_20260524"
DEFAULT_COMPARISON_PATH = ANALYSIS_DIR / "per_image_comparison.jsonl"
DEFAULT_OUTPUT_DIR = ANALYSIS_DIR / "visual_review_images"
DEFAULT_REFERENCE_PATH = (
    PROJECT_ROOT
    / "runs/validation_eval/qwen3_5_35b_a3b_exp_a_official/reference/reference_subset.jsonl"
)

DEFAULT_RUNS = {
    "A_official": PROJECT_ROOT / "runs/validation_eval/qwen3_5_35b_a3b_exp_a_official/predictions/predictions.jsonl",
    "B_full_rerun": PROJECT_ROOT
    / "runs/validation_eval/qwen3_5_35b_a3b_exp_b_attr_rerun_20260501_174649_max12000_stream/predictions/predictions.jsonl",
    "B_old_occlusion_only": PROJECT_ROOT
    / "runs/validation_eval/qwen3_5_35b_a3b_ablation_occlusion_only_tp4_20260501_2058/predictions/predictions.jsonl",
    "B_lite_dataset": PROJECT_ROOT
    / "runs/validation_eval/qwen3_5_35b_a3b_exp_b_lite_v0_20260504_2206_dataset/predictions/predictions.jsonl",
}

DEFAULT_PANEL_WIDTH = 420
DEFAULT_PANEL_HEIGHT = 300
PANEL_SPACING = 16
MARGIN = 18
TITLE_HEIGHT = 96
PANEL_HEADER_HEIGHT = 68
FOOTER_HEIGHT = 36
BACKGROUND_COLOR = (246, 244, 239)
PANEL_BG_COLOR = (232, 230, 224)
TEXT_COLOR = (28, 27, 25)
SUBTEXT_COLOR = (78, 74, 68)
GT_COLOR = (28, 138, 68)
PRED_COLOR = (210, 58, 50)
FONT_PATH = Path("/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc")
FONT_BOLD_PATH = Path("/usr/share/fonts/opentype/noto/NotoSansCJK-Bold.ttc")


def main() -> None:
    """命令行入口。"""
    parser = argparse.ArgumentParser(description="为统一错误分析导出人工复核对比图。")
    parser.add_argument("--comparison-path", default=str(DEFAULT_COMPARISON_PATH), help="per_image_comparison.jsonl")
    parser.add_argument("--reference-path", default=str(DEFAULT_REFERENCE_PATH), help="reference_subset.jsonl")
    parser.add_argument("--output-dir", default=str(DEFAULT_OUTPUT_DIR), help="图像输出目录")
    parser.add_argument("--top-k", type=int, default=30, help="导出最高风险样本数量")
    parser.add_argument("--panel-width", type=int, default=DEFAULT_PANEL_WIDTH, help="单个面板宽度")
    parser.add_argument("--panel-height", type=int, default=DEFAULT_PANEL_HEIGHT, help="单个面板高度")
    parser.add_argument("--dpi", type=int, default=150, help="保存图片 DPI")
    parser.add_argument(
        "--run",
        action="append",
        default=[],
        help="预测文件，格式为 name=predictions.jsonl。未传入时使用默认四组实验。",
    )
    args = parser.parse_args()

    runs = _parse_runs(args.run)
    comparison_rows = _read_jsonl(Path(args.comparison_path))[: args.top_k]
    reference_by_key = _load_reference_by_key(Path(args.reference_path))
    predictions_by_run = {
        run_name: _load_predictions_by_key(prediction_path)
        for run_name, prediction_path in runs.items()
    }

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    manifest_rows: list[dict[str, Any]] = []
    digits = max(3, int(math.log10(max(1, len(comparison_rows)))) + 1)
    for index, comparison in enumerate(comparison_rows, start=1):
        image_key = str(comparison["image_key"])
        reference_row = reference_by_key.get(image_key)
        if reference_row is None:
            continue
        image_path = Path(reference_row["image_path"])
        if not image_path.exists():
            continue

        output_path = output_dir / f"{index:0{digits}d}__risk{comparison['max_risk_score']}__{image_path.stem}.jpg"
        manifest_rows.append(
            _build_review_image(
                image_path=image_path,
                reference_row=reference_row,
                comparison=comparison,
                predictions_by_run=predictions_by_run,
                output_path=output_path,
                panel_width=args.panel_width,
                panel_height=args.panel_height,
                dpi=args.dpi,
            )
        )

    manifest_path = output_dir / "manifest.jsonl"
    with manifest_path.open("w", encoding="utf-8") as file:
        for row in manifest_rows:
            file.write(json.dumps(row, ensure_ascii=False))
            file.write("\n")

    summary = {
        "created_at": datetime.now().isoformat(timespec="seconds"),
        "comparison_path": str(Path(args.comparison_path).resolve()),
        "reference_path": str(Path(args.reference_path).resolve()),
        "output_dir": str(output_dir.resolve()),
        "image_count": len(manifest_rows),
        "top_k": args.top_k,
        "panel_width": args.panel_width,
        "panel_height": args.panel_height,
        "dpi": args.dpi,
        "runs": {name: str(path.resolve()) for name, path in runs.items()},
        "manifest_path": str(manifest_path.resolve()),
    }
    (output_dir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    _write_readme(output_dir / "README.md", summary, manifest_rows)
    print(json.dumps(summary, ensure_ascii=False, indent=2))


def _parse_runs(raw_runs: list[str]) -> dict[str, Path]:
    """解析运行配置。"""
    if not raw_runs:
        return DEFAULT_RUNS

    runs: dict[str, Path] = {}
    for raw_run in raw_runs:
        if "=" not in raw_run:
            raise ValueError("--run 必须使用 name=predictions.jsonl 格式。")
        name, path = raw_run.split("=", 1)
        name = name.strip()
        if not name:
            raise ValueError("--run name 不能为空。")
        runs[name] = Path(path).resolve()
    return runs


def _build_review_image(
    *,
    image_path: Path,
    reference_row: dict[str, Any],
    comparison: dict[str, Any],
    predictions_by_run: dict[str, dict[str, dict[str, Any]]],
    output_path: Path,
    panel_width: int,
    panel_height: int,
    dpi: int,
) -> dict[str, Any]:
    """生成单张复核对比图。"""
    title_font = _load_font(24, bold=True)
    header_font = _load_font(17, bold=True)
    text_font = _load_font(14)
    small_font = _load_font(12)

    panel_names = ["GT", *predictions_by_run.keys()]
    canvas_width = MARGIN * 2 + panel_width * len(panel_names) + PANEL_SPACING * (len(panel_names) - 1)
    canvas_height = TITLE_HEIGHT + PANEL_HEADER_HEIGHT + panel_height + FOOTER_HEIGHT + MARGIN

    with Image.open(image_path).convert("RGB") as image:
        canvas = Image.new("RGB", (canvas_width, canvas_height), BACKGROUND_COLOR)
        draw = ImageDraw.Draw(canvas)
        _draw_page_header(draw, comparison, image_path, title_font, text_font)

        gt_annotations = _extract_gt_annotations(reference_row)
        panels: list[tuple[str, list[dict[str, Any]], dict[str, Any] | None, tuple[int, int, int]]] = [
            ("GT", gt_annotations, None, GT_COLOR)
        ]
        for run_name, prediction_by_key in predictions_by_run.items():
            prediction_row = prediction_by_key.get(str(comparison["image_key"]))
            annotations = prediction_row.get("annotations", []) if prediction_row else []
            run_metrics = comparison.get("runs", {}).get(run_name, {})
            panels.append((run_name, annotations, run_metrics, PRED_COLOR))

        for index, (panel_name, annotations, metrics, color) in enumerate(panels):
            left = MARGIN + index * (panel_width + PANEL_SPACING)
            _draw_panel(
                canvas=canvas,
                image=image,
                left=left,
                top=TITLE_HEIGHT,
                panel_name=panel_name,
                annotations=annotations,
                metrics=metrics,
                color=color,
                panel_width=panel_width,
                panel_height=panel_height,
                header_font=header_font,
                text_font=text_font,
                small_font=small_font,
            )

    output_path.parent.mkdir(parents=True, exist_ok=True)
    canvas.save(output_path, quality=96, dpi=(dpi, dpi))
    return {
        "image_key": comparison["image_key"],
        "image_path": str(image_path.resolve()),
        "review_image": str(output_path.resolve()),
        "max_risk_score": comparison.get("max_risk_score"),
        "risk_run_count": comparison.get("risk_run_count"),
        "error_tags": comparison.get("error_tags", []),
        "best_f1_run": comparison.get("best_f1_run"),
        "worst_f1_run": comparison.get("worst_f1_run"),
    }


def _draw_page_header(
    draw: ImageDraw.ImageDraw,
    comparison: dict[str, Any],
    image_path: Path,
    title_font: ImageFont.ImageFont,
    text_font: ImageFont.ImageFont,
) -> None:
    """绘制页面标题。"""
    title = f"Risk {comparison.get('max_risk_score')} | {image_path.name}"
    tags = ", ".join(comparison.get("error_tags", [])) or "-"
    meta = (
        f"risk_runs={comparison.get('risk_run_count')} | "
        f"best_f1={comparison.get('best_f1_run')} | worst_f1={comparison.get('worst_f1_run')}"
    )
    draw.text((MARGIN, 16), title, fill=TEXT_COLOR, font=title_font)
    draw.text((MARGIN, 48), meta, fill=SUBTEXT_COLOR, font=text_font)
    draw.text((MARGIN, 70), f"tags: {tags}", fill=SUBTEXT_COLOR, font=text_font)


def _draw_panel(
    *,
    canvas: Image.Image,
    image: Image.Image,
    left: int,
    top: int,
    panel_name: str,
    annotations: list[dict[str, Any]],
    metrics: dict[str, Any] | None,
    color: tuple[int, int, int],
    panel_width: int,
    panel_height: int,
    header_font: ImageFont.ImageFont,
    text_font: ImageFont.ImageFont,
    small_font: ImageFont.ImageFont,
) -> None:
    """绘制单个面板。"""
    draw = ImageDraw.Draw(canvas)
    draw.rectangle(
        (left, top, left + panel_width, top + PANEL_HEADER_HEIGHT + panel_height),
        fill=PANEL_BG_COLOR,
    )
    draw.text((left + 10, top + 8), panel_name, fill=TEXT_COLOR, font=header_font)
    if metrics is None:
        metric_text = f"targets={len(annotations)}"
        metric_text_2 = "ground truth"
    else:
        metric_text = (
            f"GT={metrics.get('reference_count', '-')} Pred={metrics.get('prediction_count', '-')} "
            f"Match={metrics.get('matched_count', '-')}"
        )
        metric_text_2 = (
            f"P={_fmt(metrics.get('precision'))} R={_fmt(metrics.get('recall'))} "
            f"F1={_fmt(metrics.get('f1'))} M={_fmt(metrics.get('maturity_accuracy_on_matched'))} "
            f"O={_fmt(metrics.get('occlusion_accuracy_on_matched'))}"
        )
    draw.text((left + 10, top + 34), metric_text, fill=SUBTEXT_COLOR, font=small_font)
    draw.text((left + 10, top + 50), metric_text_2, fill=SUBTEXT_COLOR, font=small_font)

    preview = _fit_image(_draw_boxes(image, annotations, color=color), panel_width, panel_height)
    paste_x = left + (panel_width - preview.width) // 2
    paste_y = top + PANEL_HEADER_HEIGHT + (panel_height - preview.height) // 2
    canvas.paste(preview, (paste_x, paste_y))


def _draw_boxes(
    image: Image.Image,
    annotations: list[dict[str, Any]],
    *,
    color: tuple[int, int, int],
) -> Image.Image:
    """在图像上绘制 bbox。"""
    preview = image.copy()
    draw = ImageDraw.Draw(preview)
    label_font = _load_font(max(16, int(min(preview.size) / 110)))
    line_width = max(4, int(round(min(preview.size) / 150)))
    for index, annotation in enumerate(annotations, start=1):
        bbox = _normalize_bbox(annotation.get("bbox"))
        if bbox is None:
            continue
        bbox = _clamp_bbox(bbox, preview.width, preview.height)
        draw.rectangle(bbox, outline=color, width=line_width)
        label = _build_box_label(annotation, index)
        _draw_box_label(draw, bbox[0], bbox[1], label, color, label_font)
    return preview


def _draw_box_label(
    draw: ImageDraw.ImageDraw,
    x1: float,
    y1: float,
    text: str,
    color: tuple[int, int, int],
    font: ImageFont.ImageFont,
) -> None:
    """绘制 bbox 标签。"""
    text = text[:48]
    bbox = draw.textbbox((0, 0), text, font=font)
    padding_x = 6
    padding_y = 3
    label_height = bbox[3] - bbox[1] + 2 * padding_y
    label_width = bbox[2] - bbox[0] + 2 * padding_x
    label_box = (x1, max(0.0, y1 - label_height), x1 + label_width, max(label_height, y1))
    draw.rectangle(label_box, fill=color)
    draw.text((label_box[0] + padding_x, label_box[1] + padding_y), text, fill=(255, 255, 255), font=font)


def _build_box_label(annotation: dict[str, Any], fallback_index: int) -> str:
    """构建 bbox 标签文本。"""
    target_index = annotation.get("target_index", fallback_index - 1)
    maturity = annotation.get("maturity_level") or _label_to_maturity(annotation.get("label"))
    occlusion = annotation.get("occlusion_degree")
    parts = [f"#{target_index}"]
    if maturity:
        parts.append(str(maturity))
    if occlusion:
        parts.append(str(occlusion))
    return " ".join(parts)


def _label_to_maturity(label: Any) -> str | None:
    if not isinstance(label, str):
        return None
    for maturity in ["未成熟", "半成熟", "完熟"]:
        if maturity in label:
            return maturity
    return label[:12] if label else None


def _fit_image(image: Image.Image, max_width: int, max_height: int) -> Image.Image:
    """等比例缩放图像。"""
    scale = min(max_width / image.width, max_height / image.height)
    new_width = max(1, int(round(image.width * scale)))
    new_height = max(1, int(round(image.height * scale)))
    return image.resize((new_width, new_height), Image.Resampling.LANCZOS)


def _load_reference_by_key(path: Path) -> dict[str, dict[str, Any]]:
    """按文件名加载 reference。"""
    rows = _read_jsonl(path)
    return {Path(str(row["image_path"])).name: row for row in rows}


def _load_predictions_by_key(path: Path) -> dict[str, dict[str, Any]]:
    """按文件名加载预测结果。"""
    rows = _read_jsonl(path)
    return {Path(str(row["image_path"])).name: row for row in rows}


def _extract_gt_annotations(reference_row: dict[str, Any]) -> list[dict[str, Any]]:
    """提取 GT annotation。"""
    annotations: list[dict[str, Any]] = []
    for record in reference_row.get("source_records", []):
        if not isinstance(record, dict):
            continue
        annotations.append(
            {
                "target_index": record.get("target_index"),
                "bbox": record.get("bbox"),
                "maturity_level": record.get("maturity_level"),
                "occlusion_degree": record.get("occlusion_degree"),
            }
        )
    return annotations


def _normalize_bbox(values: Any) -> tuple[float, float, float, float] | None:
    """标准化 bbox。"""
    if not isinstance(values, list) or len(values) != 4:
        return None
    try:
        x1, y1, x2, y2 = [float(value) for value in values]
    except (TypeError, ValueError):
        return None
    if x2 <= x1 or y2 <= y1:
        return None
    return x1, y1, x2, y2


def _clamp_bbox(
    bbox: tuple[float, float, float, float],
    image_width: int,
    image_height: int,
) -> tuple[float, float, float, float]:
    """裁剪 bbox 到图像范围内。"""
    x1, y1, x2, y2 = bbox
    x1 = max(0.0, min(float(image_width - 1), x1))
    y1 = max(0.0, min(float(image_height - 1), y1))
    x2 = max(x1 + 1.0, min(float(image_width), x2))
    y2 = max(y1 + 1.0, min(float(image_height), y2))
    return x1, y1, x2, y2


def _write_readme(path: Path, summary: dict[str, Any], manifest_rows: list[dict[str, Any]]) -> None:
    """写人工查看说明。"""
    lines = [
        "# Visual Review Images",
        "",
        "本目录包含最高风险样本的五联对比图：GT、A official、B full rerun、B old + occlusion_only、B-lite dataset。",
        "",
        f"- image_count: `{summary['image_count']}`",
        f"- generated_at: `{summary['created_at']}`",
        "",
        "建议优先查看前 10 张，它们在多个实验中都被标为高风险。",
        "",
        "| # | Image Key | Max Risk | Tags | File |",
        "|---:|---|---:|---|---|",
    ]
    for index, row in enumerate(manifest_rows, start=1):
        file_name = Path(row["review_image"]).name
        tags = ", ".join(row.get("error_tags", [])) or "-"
        lines.append(f"| {index} | `{row['image_key']}` | {row['max_risk_score']} | {tags} | `{file_name}` |")
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def _read_jsonl(path: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as file:
        for line in file:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def _load_font(size: int, *, bold: bool = False) -> ImageFont.ImageFont:
    """加载支持中文的字体。"""
    font_path = FONT_BOLD_PATH if bold else FONT_PATH
    if font_path.exists():
        return ImageFont.truetype(str(font_path), size=size)
    return ImageFont.load_default()


def _fmt(value: Any) -> str:
    if value is None:
        return "-"
    try:
        return f"{float(value):.3f}"
    except (TypeError, ValueError):
        return str(value)


if __name__ == "__main__":
    main()