"""
Generate synthetic images with Stable Diffusion 3.5 Medium.

License
-------
`stabilityai/stable-diffusion-3.5-medium` is available under the Stability AI
Community License. Stage 3A conditionally approved it while Veridicate remains
under the license revenue threshold and satisfies registration requirements.
Enterprise licensing is required before use above that threshold. Do not use
outputs to create or improve a foundational generative AI model. See
`NOTICES.md` before running a real generation job.

Usage
-----
    python scripts/dataset/generate_sd35_synthetic.py \\
        --out data/raw/ai_generated/sd35-medium \\
        --count 20000 \\
        --prompts scripts/dataset/prompts.txt

Dry run, no model import/download:

    python scripts/dataset/generate_sd35_synthetic.py \\
        --out data/raw/ai_generated/sd35-medium \\
        --count 3 \\
        --dry-run

Hardware
--------
SD 3.5 Medium generation should be run on a CUDA GPU with enough VRAM for
1024x1024 half-precision inference. Start with a small smoke run on the target
GPU before renting a long job.

Idempotency
-----------
Each generated file is named by a hash of (prompt, seed), so re-runs skip
already-generated images. Crash-resume works automatically.
"""
from __future__ import annotations

import argparse
import csv
import random
from pathlib import Path

from generation_utils import (
    APPROVED_GENERATORS,
    STAGE3A_MANIFEST_FIELDS,
    choose_prompt,
    image_dimensions,
    infer_data_root,
    load_prompts,
    manifest_row,
    next_seed,
    sha256_file,
    stable_image_key,
)


FALLBACK_PROMPTS = [
    "a realistic phone photo of a grocery receipt on a kitchen table",
    "a handheld snapshot of a small dog looking out a car window",
    "a natural photo of a rainy suburban street at dusk",
    "a casual indoor photo of laundry folded on a couch",
]


SD35_SPEC = APPROVED_GENERATORS["sd3.5-medium"]


def _manifest_path(out_dir: Path) -> Path:
    return out_dir.parent / "sd35_medium_manifest.csv"


def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--out", type=Path, required=True, help="Output directory")
    parser.add_argument(
        "--data-root",
        type=Path,
        default=None,
        help=(
            "Dataset root for manifest paths. Defaults to the parent of the "
            "'raw' path segment in --out."
        ),
    )
    parser.add_argument("--count", type=int, default=20_000)
    parser.add_argument(
        "--prompts",
        type=Path,
        default=None,
        help="Optional file with one prompt per line",
    )
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--steps", type=int, default=28)
    parser.add_argument("--guidance-scale", type=float, default=4.5)
    parser.add_argument("--width", type=int, default=1024)
    parser.add_argument("--height", type=int, default=1024)
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Validate prompt/seed/output planning without loading SD 3.5 or writing files",
    )
    args = parser.parse_args()

    out_dir = args.out.resolve()
    out_dir.mkdir(parents=True, exist_ok=True)
    data_root = (
        args.data_root.resolve()
        if args.data_root is not None
        else infer_data_root(out_dir)
    )
    manifest_path = _manifest_path(out_dir)

    prompts = load_prompts(args.prompts, fallback_prompts=FALLBACK_PROMPTS)
    if args.prompts is not None:
        print(f"Loaded {len(prompts)} prompts from {args.prompts}")
    else:
        print(
            f"WARNING: --prompts not given; using {len(prompts)} built-in "
            "fallback prompts (smoke-test only, not enough diversity for "
            "a real training run)"
        )

    rng = random.Random(args.seed)

    if args.dry_run:
        print(
            "Dry run: SD 3.5 Medium pipeline will not be loaded and no images "
            "will be written."
        )
        for i in range(args.count):
            prompt = choose_prompt(rng, prompts)
            seed = next_seed(rng)
            key = stable_image_key(prompt, seed)
            dst = out_dir / f"{key}.png"
            print(f"  {i + 1:04d}: seed={seed} path={dst} prompt={prompt!r}")
        print(f"Dry run complete. Planned manifest fragment: {manifest_path}")
        return

    print("Loading SD 3.5 Medium pipeline (large download on first run)...")
    import torch
    from diffusers import StableDiffusion3Pipeline

    pipe = StableDiffusion3Pipeline.from_pretrained(
        SD35_SPEC.model_id,
        torch_dtype=torch.float16,
    )
    pipe.to("cuda")

    rows: list[dict] = []
    for i in range(args.count):
        prompt = choose_prompt(rng, prompts)
        seed = next_seed(rng)
        key = stable_image_key(prompt, seed)
        dst = out_dir / f"{key}.png"

        if not dst.exists():
            generator = torch.Generator("cuda").manual_seed(seed)
            image = pipe(
                prompt=prompt,
                num_inference_steps=args.steps,
                guidance_scale=args.guidance_scale,
                width=args.width,
                height=args.height,
                generator=generator,
            ).images[0]
            image.save(dst, format="PNG")

        width, height = image_dimensions(dst)
        rows.append(
            manifest_row(
                path=dst.relative_to(data_root).as_posix(),
                cls="ai_generated",
                spec=SD35_SPEC,
                sha256=sha256_file(dst),
                prompt=prompt,
                seed=seed,
                width=width,
                height=height,
                generation_params={
                    "steps": args.steps,
                    "guidance_scale": args.guidance_scale,
                    "width": args.width,
                    "height": args.height,
                    "pipeline": SD35_SPEC.pipeline,
                },
            )
        )

        if (i + 1) % 100 == 0:
            print(f"  generated {i + 1}/{args.count}")

    with manifest_path.open("w", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=STAGE3A_MANIFEST_FIELDS)
        writer.writeheader()
        writer.writerows(rows)

    print(f"Done. {len(rows)} images. Manifest fragment: {manifest_path}")


if __name__ == "__main__":
    main()