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

import argparse
import hashlib
import json
import os
import sys
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image, ImageDraw

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from flow_grpo.dataset_paths import DatasetPathResolver
from flow_grpo.server_profiles import apply_server_profile_defaults


apply_server_profile_defaults()
OUT_DIR = REPO_ROOT / "analysis_outputs" / "h20_eval_corruption"


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


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


def image_stats(path: Path) -> dict[str, Any]:
    image = Image.open(path)
    arr = np.asarray(image.convert("L"), dtype=np.float32)
    return {
        "path": str(path),
        "exists": True,
        "mode": image.mode,
        "size": list(image.size),
        "min": float(arr.min()),
        "max": float(arr.max()),
        "mean": float(arr.mean()),
        "std": float(arr.std()),
        "sha256": sha256_file(path),
    }


def thumb(path: Path, label: str, size: tuple[int, int] = (192, 192)) -> Image.Image:
    image = Image.open(path).convert("RGB")
    image.thumbnail((size[0], size[1] - 28), Image.Resampling.BILINEAR)
    canvas = Image.new("RGB", size, "white")
    canvas.paste(image, ((size[0] - image.width) // 2, 24 + (size[1] - 28 - image.height) // 2))
    draw = ImageDraw.Draw(canvas)
    draw.text((4, 4), label[:28], fill=(0, 0, 0))
    return canvas


def make_contact_sheet(pairs: list[dict[str, Any]], out_path: Path) -> None:
    cell_w, cell_h = 192, 192
    sheet = Image.new("RGB", (cell_w * 2, cell_h * len(pairs)), "white")
    for row, item in enumerate(pairs):
        sheet.paste(thumb(Path(item["resolved_input"]), f"{row} input"), (0, row * cell_h))
        sheet.paste(thumb(Path(item["resolved_gt"]), f"{row} gt"), (cell_w, row * cell_h))
    sheet.save(out_path)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--jsonl", default=os.environ.get("TEST_JSONL"))
    parser.add_argument("--num_samples", type=int, default=8)
    parser.add_argument("--output_dir", default=str(OUT_DIR))
    args = parser.parse_args()

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    resolver = DatasetPathResolver(os.environ.get("DATASET_ROOT"), os.environ.get("DATASET_PATH_REMAP_FROM"), os.environ.get("DATASET_PATH_REMAP_TO"))
    rows = load_rows(Path(args.jsonl).expanduser().resolve(), args.num_samples)
    audit = {
        "jsonl": args.jsonl,
        "dataset_root": str(resolver.dataset_root),
        "remap_from": resolver.remap_from,
        "remap_to": resolver.remap_to,
        "samples": [],
    }
    contact_items = []
    missing = []
    for index, sample in enumerate(rows):
        raw_input = (sample.get("input_images") or [None])[0]
        raw_gt = sample.get("output_image") or sample.get("gt_image")
        raw_mask = sample.get("output_mask") or sample.get("gt_mask") or sample.get("mask")
        resolved_input = resolver.resolve(raw_input, label=f"sample {index} input")
        resolved_gt = resolver.resolve(raw_gt, label=f"sample {index} gt")
        resolved_mask = resolver.resolve(raw_mask, label=f"sample {index} mask") if raw_mask else None
        item = {
            "index": index,
            "prompt": sample.get("prompt") or sample.get("instruction"),
            "original_input": raw_input,
            "original_gt": raw_gt,
            "original_mask": raw_mask,
            "resolved_input": str(resolved_input),
            "resolved_gt": str(resolved_gt),
            "resolved_mask": str(resolved_mask) if resolved_mask else None,
            "input_stats": image_stats(resolved_input),
            "gt_stats": image_stats(resolved_gt),
        }
        if resolved_mask and not resolved_mask.exists():
            missing.append(str(resolved_mask))
        audit["samples"].append(item)
        contact_items.append(item)

    make_contact_sheet(contact_items, out_dir / "eval_input_contact_sheet.png")
    (out_dir / "eval_input_stats.json").write_text(json.dumps(audit, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    md_lines = ["# H20 Eval Input Audit", "", f"- jsonl: `{args.jsonl}`", f"- samples: `{len(rows)}`", f"- missing: `{len(missing)}`", ""]
    for item in audit["samples"]:
        md_lines.extend([
            f"## sample {item['index']}",
            f"- input: `{item['resolved_input']}`",
            f"- gt: `{item['resolved_gt']}`",
            f"- input mean/std: `{item['input_stats']['mean']:.3f}` / `{item['input_stats']['std']:.3f}`",
            f"- gt mean/std: `{item['gt_stats']['mean']:.3f}` / `{item['gt_stats']['std']:.3f}`",
            "",
        ])
    (out_dir / "eval_input_audit.md").write_text("\n".join(md_lines), encoding="utf-8")
    print(json.dumps(audit, indent=2, sort_keys=True))
    if missing:
        raise RuntimeError(f"Missing resolved mask paths: {missing[:5]}")
    print(f"[eval-inputs] wrote outputs under {out_dir}")
    return 0


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