#!/usr/bin/env python3
"""Generate a real-photo corpus and benchmark ImageForge through its HTTP API."""

from __future__ import annotations

import argparse
import csv
import json
import math
import os
import statistics
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
from datetime import datetime, timezone
from http.cookiejar import CookieJar
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image, ImageOps


CORPUS = (
    ("photo_1015_jpeg", 1015, "JPEG", "jpg"),
    ("photo_1016_png", 1016, "PNG", "png"),
    ("photo_1025_webp", 1025, "WEBP", "webp"),
    ("photo_1039_avif", 1039, "AVIF", "avif"),
)
FORMAT_EXTENSIONS = {"jpeg": "jpg", "webp": "webp", "avif": "avif"}
PIL_FORMATS = {"jpeg": "JPEG", "webp": "WEBP", "avif": "AVIF"}


def parse_csv_arg(value: str, cast: type = str) -> list[Any]:
    return [cast(item.strip()) for item in value.split(",") if item.strip()]


def open_image(path: Path) -> Image.Image:
    with Image.open(path) as image:
        return ImageOps.exif_transpose(image).convert("RGB")


def download(url: str) -> bytes:
    request = urllib.request.Request(url, headers={"User-Agent": "ImageForge benchmark/1.0"})
    with urllib.request.urlopen(request, timeout=90) as response:
        return response.read()


def generate_corpus(output_dir: Path, width: int, height: int) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    manifest: list[dict[str, Any]] = []

    for name, image_id, image_format, extension in CORPUS:
        source_url = f"https://picsum.photos/id/{image_id}/{width}/{height}.jpg"
        source_bytes = download(source_url)
        source_path = output_dir / f".{name}.source.jpg"
        source_path.write_bytes(source_bytes)
        image = open_image(source_path)
        source_path.unlink()

        # Normalize dimensions so format and content complexity, not resolution, drive comparisons.
        image = ImageOps.fit(image, (width, height), method=Image.Resampling.LANCZOS)
        output_path = output_dir / f"{name}.{extension}"
        if image_format == "JPEG":
            image.save(output_path, format=image_format, quality=95, subsampling=0, optimize=True)
        elif image_format == "PNG":
            image.save(output_path, format=image_format, optimize=True, compress_level=9)
        elif image_format == "WEBP":
            image.save(output_path, format=image_format, lossless=True, method=6)
        else:
            image.save(output_path, format=image_format, quality=95, speed=6)

        manifest.append(
            {
                "file": output_path.name,
                "source_url": source_url,
                "picsum_id": image_id,
                "format": image_format.lower(),
                "width": image.width,
                "height": image.height,
                "size_bytes": output_path.stat().st_size,
            }
        )

    (output_dir / "manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    print(f"generated {len(manifest)} images in {output_dir}")


def multipart_body(file_path: Path, fields: dict[str, str]) -> tuple[bytes, str]:
    boundary = f"----imageforge-{uuid.uuid4().hex}"
    chunks: list[bytes] = []
    for name, value in fields.items():
        chunks.extend(
            (
                f"--{boundary}\r\n".encode(),
                f'Content-Disposition: form-data; name="{name}"\r\n\r\n'.encode(),
                value.encode(),
                b"\r\n",
            )
        )

    chunks.extend(
        (
            f"--{boundary}\r\n".encode(),
            (
                f'Content-Disposition: form-data; name="file"; '
                f'filename="{file_path.name}"\r\n'
            ).encode(),
            b"Content-Type: application/octet-stream\r\n\r\n",
            file_path.read_bytes(),
            b"\r\n",
            f"--{boundary}--\r\n".encode(),
        )
    )
    return b"".join(chunks), f"multipart/form-data; boundary={boundary}"


def request_json(
    opener: urllib.request.OpenerDirector,
    url: str,
    data: bytes,
    headers: dict[str, str],
    timeout: int,
) -> dict[str, Any]:
    request = urllib.request.Request(url, data=data, headers=headers, method="POST")
    try:
        with opener.open(request, timeout=timeout) as response:
            payload = json.load(response)
    except urllib.error.HTTPError as error:
        detail = error.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP {error.code}: {detail}") from error

    if not payload.get("success") or "data" not in payload:
        raise RuntimeError(f"unexpected API response: {payload}")
    return payload["data"]


def login(
    opener: urllib.request.OpenerDirector,
    base_url: str,
    email: str,
    password: str,
    timeout: int,
) -> str:
    body = json.dumps({"email": email, "password": password}).encode()
    data = request_json(
        opener,
        f"{base_url}/api/v1/auth/login",
        body,
        {"Content-Type": "application/json"},
        timeout,
    )
    return str(data["token"])


def block_ssim(reference: np.ndarray, candidate: np.ndarray, block: int = 8) -> float:
    ref = 0.2126 * reference[..., 0] + 0.7152 * reference[..., 1] + 0.0722 * reference[..., 2]
    out = 0.2126 * candidate[..., 0] + 0.7152 * candidate[..., 1] + 0.0722 * candidate[..., 2]
    height = (ref.shape[0] // block) * block
    width = (ref.shape[1] // block) * block
    if height == 0 or width == 0:
        height, width, block = ref.shape[0], ref.shape[1], 1

    def blocks(array: np.ndarray) -> np.ndarray:
        return (
            array[:height, :width]
            .reshape(height // block, block, width // block, block)
            .transpose(0, 2, 1, 3)
        )

    ref_blocks = blocks(ref)
    out_blocks = blocks(out)
    axes = (-1, -2)
    ref_mean = ref_blocks.mean(axis=axes)
    out_mean = out_blocks.mean(axis=axes)
    ref_var = ref_blocks.var(axis=axes)
    out_var = out_blocks.var(axis=axes)
    covariance = ((ref_blocks - ref_mean[..., None, None]) * (out_blocks - out_mean[..., None, None])).mean(axis=axes)
    c1 = (0.01 * 255.0) ** 2
    c2 = (0.03 * 255.0) ** 2
    numerator = (2 * ref_mean * out_mean + c1) * (2 * covariance + c2)
    denominator = (ref_mean**2 + out_mean**2 + c1) * (ref_var + out_var + c2)
    return float(np.mean(numerator / np.maximum(denominator, 1e-12)))


def image_metrics(reference_path: Path, output_path: Path) -> dict[str, Any]:
    reference = open_image(reference_path)
    with Image.open(output_path) as opened:
        detected_format = (opened.format or "unknown").upper()
        output = ImageOps.exif_transpose(opened).convert("RGB")

    output_width, output_height = output.size
    if output.size != reference.size:
        output = output.resize(reference.size, Image.Resampling.LANCZOS)

    ref_array = np.asarray(reference, dtype=np.float64)
    out_array = np.asarray(output, dtype=np.float64)
    mse = float(np.mean((ref_array - out_array) ** 2))
    psnr = 99.0 if mse == 0 else 20.0 * math.log10(255.0 / math.sqrt(mse))
    return {
        "detected_format": detected_format,
        "width": output_width,
        "height": output_height,
        "pixel_ratio_pct": output_width * output_height * 100.0 / (reference.width * reference.height),
        "ssim": block_ssim(ref_array, out_array),
        "psnr_db": psnr,
    }


def summarize(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    groups: dict[tuple[str, int], list[dict[str, Any]]] = {}
    for row in rows:
        if row.get("error"):
            continue
        groups.setdefault((str(row["output_format"]), int(row["requested_rate"])), []).append(row)

    summary: list[dict[str, Any]] = []
    for (output_format, requested_rate), group in sorted(groups.items()):
        summary.append(
            {
                "output_format": output_format,
                "requested_rate": requested_rate,
                "cases": len(group),
                "target_met": sum(bool(row["target_met"]) for row in group),
                "format_ok": sum(bool(row["format_ok"]) for row in group),
                "mean_actual_rate_pct": statistics.mean(float(row["actual_rate_pct"]) for row in group),
                "mean_saved_pct": statistics.mean(float(row["saved_pct"]) for row in group),
                "mean_ssim": statistics.mean(float(row["ssim"]) for row in group),
                "mean_psnr_db": statistics.mean(float(row["psnr_db"]) for row in group),
                "mean_pixel_ratio_pct": statistics.mean(float(row["pixel_ratio_pct"]) for row in group),
                "median_elapsed_ms": statistics.median(float(row["elapsed_ms"]) for row in group),
            }
        )
    return summary


def benchmark(args: argparse.Namespace) -> int:
    input_dir = Path(args.input_dir)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    inputs = sorted(path for path in input_dir.iterdir() if path.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp", ".avif"})
    if not inputs:
        raise RuntimeError(f"no benchmark images found in {input_dir}")

    rates = parse_csv_arg(args.rates, int)
    formats = [str(value).lower() for value in parse_csv_arg(args.formats)]
    unsupported = sorted(set(formats) - set(FORMAT_EXTENSIONS))
    if unsupported:
        raise RuntimeError(f"unsupported output formats: {', '.join(unsupported)}")

    base_url = args.base_url.rstrip("/")
    opener = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(CookieJar()))
    token = args.token or os.getenv("IMAGEFORGE_BENCH_TOKEN", "")
    if not token:
        email = args.email or os.getenv("IMAGEFORGE_BENCH_EMAIL", "")
        password = args.password or os.getenv("IMAGEFORGE_BENCH_PASSWORD", "")
        if email and password:
            token = login(opener, base_url, email, password, args.timeout)

    auth_headers = {"Authorization": f"Bearer {token}"} if token else {}
    rows: list[dict[str, Any]] = []
    total = len(inputs) * len(formats) * len(rates)
    case_number = 0

    for input_path in inputs:
        original_size = input_path.stat().st_size
        for output_format in formats:
            for rate in rates:
                case_number += 1
                output_path = output_dir / f"{input_path.stem}__{output_format}__r{rate}.{FORMAT_EXTENSIONS[output_format]}"
                row: dict[str, Any] = {
                    "input": input_path.name,
                    "input_format": input_path.suffix.lower().lstrip("."),
                    "output_format": output_format,
                    "requested_rate": rate,
                    "original_size": original_size,
                }
                try:
                    body, content_type = multipart_body(
                        input_path,
                        {"compression_rate": str(rate), "output_format": output_format},
                    )
                    started = time.perf_counter()
                    data = request_json(
                        opener,
                        f"{base_url}/api/v1/compress",
                        body,
                        {**auth_headers, "Content-Type": content_type},
                        args.timeout,
                    )
                    elapsed_ms = (time.perf_counter() - started) * 1000.0

                    download_request = urllib.request.Request(
                        urllib.parse.urljoin(f"{base_url}/", str(data["download_url"]).lstrip("/")),
                        headers=auth_headers,
                    )
                    with opener.open(download_request, timeout=args.timeout) as response:
                        output_path.write_bytes(response.read())

                    compressed_size = output_path.stat().st_size
                    metrics = image_metrics(input_path, output_path)
                    tolerance_bytes = max(1024, int(original_size * 0.01))
                    target_bytes = original_size * rate / 100.0
                    row.update(
                        {
                            "compressed_size": compressed_size,
                            "actual_rate_pct": compressed_size * 100.0 / original_size,
                            "saved_pct": max(0.0, (original_size - compressed_size) * 100.0 / original_size),
                            "target_error_pct_points": compressed_size * 100.0 / original_size - rate,
                            "target_met": compressed_size <= target_bytes + tolerance_bytes,
                            "format_ok": metrics["detected_format"] == PIL_FORMATS[output_format],
                            "api_size_matches": int(data["compressed_size"]) == compressed_size,
                            "elapsed_ms": elapsed_ms,
                            **metrics,
                            "error": "",
                        }
                    )
                except Exception as error:  # Continue to expose the full failure matrix.
                    row["error"] = str(error)
                rows.append(row)
                status = "ERROR" if row.get("error") else f"{row['actual_rate_pct']:.1f}% SSIM={row['ssim']:.4f}"
                print(f"[{case_number:02d}/{total:02d}] {input_path.name} -> {output_format} r{rate}: {status}", flush=True)
                if args.delay:
                    time.sleep(args.delay)

    summary = summarize(rows)
    report = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "base_url": base_url,
        "inputs": len(inputs),
        "cases": len(rows),
        "errors": sum(bool(row.get("error")) for row in rows),
        "format_failures": sum(not bool(row.get("format_ok")) for row in rows if not row.get("error")),
        "target_failures": sum(not bool(row.get("target_met")) for row in rows if not row.get("error")),
        "summary": summary,
        "results": rows,
    }
    (output_dir / "report.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    fieldnames = sorted({key for row in rows for key in row})
    with (output_dir / "results.csv").open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    print("\nformat rate target format-ok actual% saved% SSIM PSNR pixel% median-ms")
    for item in summary:
        print(
            f"{item['output_format']:>6} {item['requested_rate']:>4} "
            f"{item['target_met']}/{item['cases']} {item['format_ok']}/{item['cases']} "
            f"{item['mean_actual_rate_pct']:>7.2f} {item['mean_saved_pct']:>6.2f} "
            f"{item['mean_ssim']:.4f} {item['mean_psnr_db']:>5.2f} "
            f"{item['mean_pixel_ratio_pct']:>6.2f} {item['median_elapsed_ms']:>9.1f}"
        )
    print(f"\nreport: {output_dir / 'report.json'}")
    return 1 if args.strict and (report["errors"] or report["format_failures"] or report["target_failures"]) else 0


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)

    generate = subparsers.add_parser("generate", help="download and build the fixed real-photo corpus")
    generate.add_argument("--output-dir", default=".bench/corpus")
    generate.add_argument("--width", type=int, default=960)
    generate.add_argument("--height", type=int, default=640)

    run = subparsers.add_parser("run", help="benchmark a running ImageForge deployment")
    run.add_argument("--base-url", default="http://127.0.0.1:8080")
    run.add_argument("--input-dir", default=".bench/corpus")
    run.add_argument("--output-dir", default=".bench/results")
    run.add_argument("--formats", default="jpeg,webp,avif")
    run.add_argument("--rates", default="30,50,70")
    run.add_argument("--email", default="")
    run.add_argument("--password", default="")
    run.add_argument("--token", default="")
    run.add_argument("--timeout", type=int, default=300)
    run.add_argument("--delay", type=float, default=0.05)
    run.add_argument("--strict", action="store_true")
    return parser


def main() -> int:
    args = build_parser().parse_args()
    if args.command == "generate":
        generate_corpus(Path(args.output_dir), args.width, args.height)
        return 0
    return benchmark(args)


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except KeyboardInterrupt:
        raise SystemExit(130) from None
    except Exception as error:
        print(f"benchmark failed: {error}", file=sys.stderr)
        raise SystemExit(1) from error