"""
3_generate.py  —  STAGE 3: generate images with your trained LoRA
=================================================================
Loads FLUX.1-dev, applies the LoRA you trained in Stage 2, and generates images
from prompts that include your trigger word. This is the LoRA-aware version of
your app.py's generate path (minus the S3/Mongo/webhook plumbing, which you can
add back from app.py once generation works).

Run (single prompt):
    python 3_generate.py --lora ./flux-lora-out/gr4v_style_step1000 \
        --prompt "gr4v_style a cinematic portrait, golden hour" --out result.png

Run (batch from a file, one prompt per line):
    python 3_generate.py --lora ./flux-lora-out/gr4v_style_step1000 \
        --prompts prompts.txt --outdir ./generated

Install:
    pip install -U "diffusers>=0.32" transformers peft safetensors torch pillow
"""

import argparse
from pathlib import Path

import torch
from diffusers import FluxPipeline


def build_pipe(base_model: str, lora_path: str, weight: float, device):
    pipe = FluxPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
    pipe.to(device)

    # Load the trained LoRA. lora_path is the folder saved by Stage 2.
    pipe.load_lora_weights(lora_path)
    # Control LoRA strength (0.7-1.0 typical). Higher = stronger subject.
    try:
        pipe.set_adapters(["default_0"], adapter_weights=[weight])
    except Exception:
        pass  # older diffusers apply at full strength by default
    return pipe


def generate_one(pipe, prompt, out_path, steps, guidance, height, width, seed, device):
    g = torch.Generator(device).manual_seed(seed)
    image = pipe(
        prompt=prompt,
        num_inference_steps=steps,
        guidance_scale=guidance,
        height=height, width=width,
        generator=g,
    ).images[0]
    image.save(out_path)
    print(f"  saved -> {out_path}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--lora", required=True, help="folder with trained LoRA weights")
    ap.add_argument("--base", default="black-forest-labs/FLUX.1-dev")
    ap.add_argument("--prompt", help="single prompt (include your trigger word)")
    ap.add_argument("--prompts", help="text file, one prompt per line")
    ap.add_argument("--out", default="result.png", help="output for --prompt")
    ap.add_argument("--outdir", default="./generated", help="output dir for --prompts")
    ap.add_argument("--weight", type=float, default=0.9, help="LoRA strength 0.0-1.0")
    ap.add_argument("--steps", type=int, default=28)
    ap.add_argument("--guidance", type=float, default=3.5)
    ap.add_argument("--height", type=int, default=1024)
    ap.add_argument("--width", type=int, default=1024)
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print("Loading FLUX + LoRA...")
    pipe = build_pipe(args.base, args.lora, args.weight, device)

    if args.prompts:
        outdir = Path(args.outdir)
        outdir.mkdir(parents=True, exist_ok=True)
        lines = [l.strip() for l in Path(args.prompts).read_text().splitlines() if l.strip()]
        for i, prompt in enumerate(lines):
            generate_one(pipe, prompt, outdir / f"gen_{i:03d}.png",
                         args.steps, args.guidance, args.height, args.width,
                         args.seed + i, device)
    elif args.prompt:
        generate_one(pipe, args.prompt, args.out,
                     args.steps, args.guidance, args.height, args.width,
                     args.seed, device)
    else:
        raise SystemExit("Provide --prompt '...' or --prompts file.txt")

    print("Done.")


if __name__ == "__main__":
    main()