#!/usr/bin/env python3
"""Generate one conventional 50-step Clover Image Tiny preview image."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
from diffusers import DiffusionPipeline, PNDMScheduler


def _runtime(requested: str) -> tuple[torch.device, torch.dtype]:
    if requested == "auto":
        if torch.cuda.is_available():
            device = torch.device("cuda")
        elif torch.backends.mps.is_available():
            device = torch.device("mps")
        else:
            device = torch.device("cpu")
    else:
        device = torch.device(requested)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA was requested but is unavailable")
    if device.type == "mps" and not torch.backends.mps.is_available():
        raise RuntimeError("MPS was requested but is unavailable")
    dtype = torch.float16 if device.type in {"cuda", "mps"} else torch.float32
    return device, dtype


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Run Clover Image Tiny with its fixed conventional preview settings."
    )
    parser.add_argument("--model", required=True, help="Hub repository ID or local directory")
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--seed", type=int, default=1337)
    parser.add_argument("--output", type=Path, default=Path("clover-image-tiny.png"))
    parser.add_argument("--device", choices=("auto", "cuda", "mps", "cpu"), default="auto")
    parser.add_argument("--local-files-only", action="store_true")
    return parser


def main() -> int:
    args = _parser().parse_args()
    device, dtype = _runtime(args.device)
    pipe = DiffusionPipeline.from_pretrained(
        args.model,
        torch_dtype=dtype,
        local_files_only=args.local_files_only,
    )
    pipe.scheduler = PNDMScheduler.from_config(pipe.scheduler.config)
    pipe = pipe.to(device)
    generator_device = "cuda" if device.type == "cuda" else "cpu"
    generator = torch.Generator(device=generator_device).manual_seed(args.seed)
    with torch.inference_mode():
        result = pipe(
            prompt=args.prompt,
            negative_prompt="",
            num_inference_steps=50,
            guidance_scale=7.5,
            height=512,
            width=512,
            generator=generator,
        )
    if not result.images:
        raise RuntimeError("the pipeline returned no images")
    args.output.parent.mkdir(parents=True, exist_ok=True)
    result.images[0].save(args.output)
    metadata = {
        "model": args.model,
        "prompt": args.prompt,
        "negative_prompt": "",
        "seed": args.seed,
        "scheduler": type(pipe.scheduler).__name__,
        "num_inference_steps": 50,
        "guidance_scale": 7.5,
        "height": 512,
        "width": 512,
        "device": device.type,
        "dtype": str(dtype).removeprefix("torch."),
        "nsfw_content_detected": getattr(result, "nsfw_content_detected", None),
        "cross_device_pixel_identity_claimed": False,
        "usage_label": "RESEARCH PREVIEW - PRODUCTION PILOT - NOT RELEASE-READY",
    }
    metadata_path = args.output.with_suffix(args.output.suffix + ".json")
    metadata_path.write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")
    print(json.dumps({"output": str(args.output), "metadata": str(metadata_path), **metadata}))
    return 0


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