"""
Generate synthetic images with Flux.1-schnell.

Flux.1-schnell is released by Black Forest Labs under Apache 2.0. Outputs are
not restricted, so we can use them as commercial training data.

Usage
-----
    python scripts/dataset/generate_flux_synthetic.py \
        --out data/raw/ai_generated \
        --count 50000 \
        --prompts scripts/dataset/prompts.txt

Prompts file format
-------------------
One complete prompt per line. Blank lines and lines starting with `#` are
ignored. Diversity is critical: a model trained on monotonous prompts learns
to detect prompt style, not AI-generation artefacts. Aim for the prompts
file to span the full distribution of subjects, settings, lighting, and
camera styles a regular person might upload — including mundane and
imperfect ones (snapshots, blurry shots, boring objects), not just
gallery-worthy compositions.

If `--prompts` is omitted, a tiny built-in fallback set is used. That set
is only sufficient for smoke tests, not for a real training run.

Hardware
--------
Flux.1-schnell needs a GPU with ~12-16 GB VRAM. Runs comfortably on:
  • RunPod / Lambda Labs / Vast.ai spot instances (RTX 3090 / 4090 / A10G)
  • ~3-5 seconds per image at 4 inference steps on an A10G
  • 50k images ≈ 60-80 GPU-hours ≈ $25-50 at typical spot rates

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,
    choose_prompt,
    image_dimensions,
    load_prompts,
    manifest_row,
    next_seed,
    sha256_file,
    stable_image_key,
)


# Fallback used only when --prompts is not provided. Intentionally tiny —
# enough for a smoke test, not for a real training run.
FALLBACK_PROMPTS = [
    "a photo of a dog in a park, golden hour",
    "candid photograph of a person walking on a busy street",
    "studio lighting portrait of two friends talking",
    "a wide-angle shot of a kitchen with a child playing",
]


FLUX_SPEC = APPROVED_GENERATORS["flux.1-schnell"]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path, required=True, help="Output directory")
    parser.add_argument("--count", type=int, default=50_000)
    parser.add_argument(
        "--prompts",
        type=Path,
        default=None,
        help="Optional file with one prompt template per line",
    )
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--steps", type=int, default=4, help="Flux schnell uses 4 steps")
    args = parser.parse_args()

    args.out.mkdir(parents=True, exist_ok=True)
    manifest_path = args.out.parent / "ai_generated_manifest.csv"

    if args.prompts is not None:
        prompts = load_prompts(args.prompts)
        print(f"Loaded {len(prompts)} prompts from {args.prompts}")
    else:
        prompts = load_prompts(None, fallback_prompts=FALLBACK_PROMPTS)
        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)

    # Lazy import — diffusers + torch + accelerate are heavy.
    print("Loading Flux.1-schnell pipeline (~24 GB download on first run)...")
    import torch
    from diffusers import FluxPipeline

    pipe = FluxPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-schnell",
        torch_dtype=torch.bfloat16,
    )
    pipe.to("cuda")

    rows: list[dict] = []
    for i in range(args.count):
        prompt = choose_prompt(rng, prompts)
        seed = next_seed(rng)
        # Filename is content-addressable so re-runs are idempotent.
        key = stable_image_key(prompt, seed)
        dst = args.out / f"{key}.png"

        if not dst.exists():
            generator = torch.Generator("cuda").manual_seed(seed)
            image = pipe(
                prompt,
                guidance_scale=0.0,  # Flux schnell ignores guidance
                num_inference_steps=args.steps,
                generator=generator,
            ).images[0]
            image.save(dst, format="PNG")

        width, height = image_dimensions(dst)

        rows.append(manifest_row(
            # as_posix() — keep manifests portable across OS boundaries.
            path=dst.relative_to(args.out.parent.parent).as_posix(),
            cls="ai_generated",
            spec=FLUX_SPEC,
            sha256=sha256_file(dst),
            prompt=prompt,
            seed=seed,
            width=width,
            height=height,
            generation_params={
                "steps": args.steps,
                "guidance_scale": 0.0,
                "pipeline": FLUX_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=[
                "path", "class", "source", "license", "license_url",
                "sha256", "generator", "model_family", "model_id", "prompt",
                "seed", "width", "height", "generation_params_json",
            ],
        )
        writer.writeheader()
        writer.writerows(rows)

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


if __name__ == "__main__":
    main()