#!/usr/bin/env python3
"""Z-Image CLI - Generate images from text prompts"""

import argparse
import sys
from pathlib import Path
from datetime import datetime

from diffusers import ZImagePipeline
import torch


def main():
    parser = argparse.ArgumentParser(description="Generate images with Z-Image")
    parser.add_argument("prompt", help="Text description of the image to generate")
    parser.add_argument("-o", "--output", help="Output filename (default: auto-generated)")
    parser.add_argument("-s", "--steps", type=int, default=9, help="Inference steps (default: 9)")
    parser.add_argument("--no-offload", action="store_true", help="Disable CPU offloading (requires more VRAM)")
    parser.add_argument("-d", "--output-dir", default=".", help="Output directory (default: current)")

    args = parser.parse_args()

    print("Loading Z-Image model...")
    pipe = ZImagePipeline.from_pretrained(
        "Tongyi-MAI/Z-Image-Turbo",
        torch_dtype=torch.bfloat16
    )

    if args.no_offload:
        pipe.to("cuda")
    else:
        pipe.enable_model_cpu_offload()

    print(f"Generating image with {args.steps} steps...")
    print(f"Prompt: {args.prompt}")

    image = pipe(
        prompt=args.prompt,
        num_inference_steps=args.steps,
        guidance_scale=0.0
    ).images[0]

    # Generate output filename if not provided
    if args.output:
        output_path = Path(args.output_dir) / args.output
    else:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        # Create safe filename from prompt
        safe_prompt = "".join(c if c.isalnum() or c in " -_" else "" for c in args.prompt[:30])
        safe_prompt = safe_prompt.strip().replace(" ", "_")
        output_path = Path(args.output_dir) / f"zimage_{timestamp}_{safe_prompt}.png"

    output_path.parent.mkdir(parents=True, exist_ok=True)
    image.save(output_path)
    print(f"Image saved to: {output_path}")

    return str(output_path)


if __name__ == "__main__":
    main()