"""Create a deterministic, held-out synthetic restoration benchmark.

Only real-portrait folders are selected. The generated manifest and holdout
list are ignored by Git; pass the holdout list to training to avoid leakage.
"""
from __future__ import annotations

import argparse
import csv
import hashlib
import shutil
from pathlib import Path

import cv2
import numpy as np

ROOT = Path(__file__).resolve().parents[1]
DEFAULT_SOURCE = ROOT / "models" / "CodeFormer" / "datasets" / "ffhq" / "ffhq_512"
REAL_FOLDERS = {"faces", "pinterest"}
EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp"}


def candidates(source: Path) -> list[Path]:
    files = []
    for folder in REAL_FOLDERS:
        directory = source / folder
        if directory.is_dir():
            files.extend(p for p in directory.rglob("*") if p.suffix.lower() in EXTENSIONS)
    return sorted(files)


def degrade(image: np.ndarray, seed: int) -> np.ndarray:
    rng = np.random.default_rng(seed)
    blur_size = int(rng.choice([3, 5, 7]))
    output = cv2.GaussianBlur(image, (blur_size, blur_size), float(rng.uniform(0.4, 2.5)))
    h, w = output.shape[:2]
    scale = float(rng.uniform(1.5, 6.0))
    small = cv2.resize(output, (max(32, int(w / scale)), max(32, int(h / scale))), interpolation=cv2.INTER_AREA)
    output = cv2.resize(small, (w, h), interpolation=cv2.INTER_CUBIC)
    output = np.clip(output.astype(np.float32) + rng.normal(0, rng.uniform(2, 25), output.shape), 0, 255).astype(np.uint8)
    quality = int(rng.integers(10, 71))
    ok, encoded = cv2.imencode(".jpg", output, [cv2.IMWRITE_JPEG_QUALITY, quality])
    if not ok:
        raise RuntimeError("JPEG degradation failed")
    return cv2.imdecode(encoded, cv2.IMREAD_COLOR)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source", type=Path, default=DEFAULT_SOURCE)
    parser.add_argument("--count", type=int, default=500)
    parser.add_argument("--seed", type=int, default=20260723)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    source = args.source.resolve()
    ranked = sorted(candidates(source), key=lambda p: hashlib.sha256(f"{args.seed}:{p.relative_to(source)}".encode()).hexdigest())
    selected = ranked[:args.count]
    if len(selected) < args.count:
        raise SystemExit(f"Need {args.count} real portraits but found {len(selected)} in {source}")
    print(f"Selected {len(selected)} held-out real portraits from {len(ranked)} candidates.")
    if args.dry_run:
        return
    benchmark = ROOT / "benchmarks"
    inputs, references = benchmark / "inputs", benchmark / "references"
    inputs.mkdir(parents=True, exist_ok=True); references.mkdir(parents=True, exist_ok=True)
    manifest = benchmark / "manifest.csv"
    holdout = benchmark / "holdout_paths.txt"
    with manifest.open("w", newline="", encoding="utf-8") as csv_file, holdout.open("w", encoding="utf-8") as holdout_file:
        writer = csv.DictWriter(csv_file, fieldnames=["id", "input_path", "reference_path", "category", "notes"]); writer.writeheader()
        for index, path in enumerate(selected, 1):
            image = cv2.imread(str(path), cv2.IMREAD_COLOR)
            if image is None: raise RuntimeError(f"Unreadable image: {path}")
            sample_id = f"portrait_{index:04d}"
            cv2.imwrite(str(inputs / f"{sample_id}.jpg"), degrade(image, args.seed + index), [cv2.IMWRITE_JPEG_QUALITY, 95])
            shutil.copy2(path, references / f"{sample_id}{path.suffix.lower()}")
            writer.writerow({"id": sample_id, "input_path": f"inputs/{sample_id}.jpg", "reference_path": f"references/{sample_id}{path.suffix.lower()}", "category": "synthetic_real_portrait", "notes": "deterministic holdout"})
            holdout_file.write(path.relative_to(source).as_posix() + "\n")
    print(f"Created {manifest} and {holdout}. Do not train on the listed paths.")

if __name__ == "__main__": main()