#!/usr/bin/env python3
"""Exercise package file contracts; optionally execute one real model forward."""

from __future__ import annotations

import argparse
import subprocess
import sys
import tempfile
from pathlib import Path

DELIVERY_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(DELIVERY_ROOT))

import numpy as np
from PIL import Image

from benchmark.validate_outputs import validate_outputs
from raesr.artifacts import verify_decoder, verify_frontend
from scripts.infer import collect_images, is_valid_output


def make_image(path: Path, size: tuple[int, int]) -> None:
    y, x = np.indices((size[1], size[0]))
    array = np.stack(((x * 3) % 256, (y * 5) % 256, ((x + y) * 7) % 256), axis=-1).astype(np.uint8)
    Image.fromarray(array).save(path)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--decoder", type=Path)
    parser.add_argument("--frontend", type=Path)
    parser.add_argument("--model", choices=("raesr", "raesr-d"), default="raesr")
    parser.add_argument("--dinov3-root", type=Path)
    parser.add_argument("--dinov3-weights", type=Path)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--with-model", action="store_true")
    args = parser.parse_args()
    if args.with_model and (args.decoder is None or args.dinov3_root is None or args.dinov3_weights is None):
        parser.error("--with-model requires --decoder, --dinov3-root, and --dinov3-weights")
    if args.with_model and args.model == "raesr-d" and args.frontend is None:
        parser.error("--with-model --model raesr-d requires --frontend")

    with tempfile.TemporaryDirectory(prefix="raesr-smoke-") as temporary:
        root = Path(temporary)
        inputs, outputs = root / "inputs", root / "outputs"
        inputs.mkdir()
        outputs.mkdir()
        make_image(inputs / "fixture.png", (63, 47))
        assert len(collect_images(inputs)) == 1
        make_image(outputs / "fixture.png", (512, 512))
        assert is_valid_output(outputs / "fixture.png")
        validate_outputs(inputs, outputs, expected_count=1)
        if args.decoder:
            verify_decoder(args.decoder)
        if args.frontend:
            verify_frontend(args.frontend)
        if args.with_model:
            outputs.mkdir(exist_ok=True)
            command = [
                sys.executable, str(DELIVERY_ROOT / "scripts" / "infer.py"),
                "--input-dir", str(inputs), "--output-dir", str(outputs), "--model", args.model,
                "--decoder", str(args.decoder),
                "--dinov3-root", str(args.dinov3_root), "--dinov3-weights", str(args.dinov3_weights),
                "--device", args.device, "--overwrite",
            ]
            if args.frontend is not None:
                command += ["--frontend", str(args.frontend)]
            subprocess.run(command, check=True)
            assert is_valid_output(outputs / "fixture.png")
    print("smoke test passed")


if __name__ == "__main__":
    main()