"""
SDXS-2B image generation from prompts.txt.
Usage:
  python generate.py                          # generates from prompts.txt
  python generate.py --prompts my_list.txt    # custom file
  python generate.py --prompt "your text"     # single prompt
"""
import argparse
import os
import sys
import math
import textwrap

import torch
import numpy as np

from pipeline_sdxs import SdxsPipeline

NEGATIVE_PROMPT = (
    "low quality, bad quality, blurry, sketch, sepia, text, "
    "bad anatomy, bad proportions, bad hands, missing fingers, child drawing"
)

DEFAULT_PROMPTS_FILE = "prompts.txt"
MEDIA_DIR = "media"

def main():
    parser = argparse.ArgumentParser(description="SDXS-2B generation")
    parser.add_argument("--prompts", type=str, default=None, help="prompts file")
    parser.add_argument("--prompt", type=str, default=None, help="single prompt")
    parser.add_argument("--steps", type=int, default=40)
    parser.add_argument("--guidance", type=float, default=4.0)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--width", type=int, default=832)
    parser.add_argument("--height", type=int, default=1152)
    parser.add_argument("--output", type=str, default=None,
                        help="output dir for individual images (optional)")
    args = parser.parse_args()

    # collect prompts
    if args.prompt:
        prompts = [args.prompt]
    elif args.prompts:
        with open(args.prompts) as f:
            prompts = [l.strip() for l in f if l.strip()]
    else:
        with open(DEFAULT_PROMPTS_FILE) as f:
            prompts = [l.strip() for l in f if l.strip()]

    if not prompts:
        print("No prompts found")
        return

    # load pipeline
    pipe = SdxsPipeline.from_pretrained(
        os.path.dirname(os.path.abspath(__file__)),
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
    ).to("cuda:0")

    images = []
    for i, prompt in enumerate(prompts):
        print(f"[{i+1}/{len(prompts)}] {prompt[:80]}...")
        image = pipe(
            prompt=prompt,
            negative_prompt=NEGATIVE_PROMPT,
            guidance_scale=args.guidance,
            width=args.width,
            height=args.height,
            seed=args.seed,
            num_inference_steps=args.steps,
            show_progress_bar=False,
        )[0][0]
        images.append(image)

        if args.output:
            os.makedirs(args.output, exist_ok=True)
            path = os.path.join(args.output, f"{i+1:03d}.jpg")
            image.save(path, quality=95)
            print(f"  saved {path}")

    # save grid
    os.makedirs(MEDIA_DIR, exist_ok=True)
    grid_path = os.path.join(MEDIA_DIR, "result_grid.jpg")

    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    cols = min(4, len(images))
    rows = math.ceil(len(images) / cols)
    fig, axes = plt.subplots(rows, cols, figsize=(cols * 4, rows * 4.5), constrained_layout=True)
    axes = list(np.array(axes).flatten())[:len(images)]
    for i, (img, prompt) in enumerate(zip(images, prompts)):
        ax = axes[i]
        ax.imshow(img)
        ax.axis("off")
        ax.set_aspect("equal")
        text = (prompt[:200] + "…") if len(prompt) > 200 else prompt
        lines = textwrap.wrap(text, width=35)
        while len(lines) < 4:
            lines.append("")
        ax.set_title("\n".join(lines), fontsize=9, pad=8)
    plt.savefig(grid_path, bbox_inches="tight", dpi=150, format="jpeg")
    print(f"grid -> {grid_path}")


if __name__ == "__main__":
    main()