# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import argparse
import os
import re
from typing import List


DEFAULT_PROMPTS = [
    "a photo of a cat",
    "a cinematic photo of a red panda wearing a tiny backpack, standing on a rainy neon-lit street at night, shallow depth of field, sharp focus, 35mm, bokeh",
]


def _safe_filename(text: str, max_len: int = 100) -> str:
    """
    Make a stable, filesystem-friendly filename base.
    VideoGenerator uses prompt[:100].strip() internally, so we mirror that,
    but also remove path separators and other problematic characters.
    """
    s = text[:max_len].strip()
    s = s.replace(os.sep, "_")
    if os.altsep:
        s = s.replace(os.altsep, "_")
    s = re.sub(r"\s+", " ", s)
    s = re.sub(r"[^A-Za-z0-9 .,_-]", "_", s)
    s = s.strip(" .")
    return s or "prompt"


def _remove_existing_outputs(out_dir: str, filename_base: str) -> None:
    """
    Ensure deterministic naming by deleting any existing outputs that would
    cause VideoGenerator to append suffixes like _1, _2, etc.
    """
    if not os.path.isdir(out_dir):
        return

    pattern = re.compile(rf"^{re.escape(filename_base)}(_\d+)?\.(mp4|png)$")
    for fn in os.listdir(out_dir):
        if pattern.match(fn):
            try:
                os.remove(os.path.join(out_dir, fn))
            except FileNotFoundError:
                pass


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Run SD3.5 Medium text-to-image with FastVideo VideoGenerator.")
    p.add_argument("--model-path", default="stabilityai/stable-diffusion-3.5-medium", help="Path to local diffusers-format SD3.5 weights directory.")
    p.add_argument(
        "--out-dir",
        "--outdir",
        default="outputs/sd35/samples",
        help="Output directory for generated mp4 files.",
    )
    p.add_argument(
        "--prompt",
        action="append",
        default=None,
        help="Prompt text. Repeat --prompt multiple times to generate multiple samples.",
    )
    p.add_argument("--negative", default="lowres, blurry, jpeg artifacts, watermark, text", help="Negative prompt.")
    p.add_argument(
        "--backend",
        default=None,
        help="Set FASTVIDEO_ATTENTION_BACKEND (e.g. TORCH_SDPA). If omitted, respects the existing env var.",
    )
    p.add_argument("--seed", type=int, default=42, help="Base seed. Each prompt uses seed + prompt_idx.")
    p.add_argument("--height", type=int, default=768, help="Output height.")
    p.add_argument("--width", type=int, default=768, help="Output width.")
    p.add_argument("--steps", type=int, default=28, help="Number of inference steps.")
    p.add_argument("--guidance", type=float, default=6.0, help="Guidance scale.")
    p.add_argument("--num-gpus", type=int, default=1, help="Number of GPUs to use.")
    return p.parse_args()


def main() -> None:
    args = parse_args()
    prompts: List[str] = args.prompt if args.prompt else DEFAULT_PROMPTS

    if args.backend:
        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = args.backend

    from fastvideo import VideoGenerator
    from fastvideo.api import (
        EngineConfig, GenerationRequest, GeneratorConfig, OffloadConfig, OutputConfig,
        ParallelismConfig, PipelineSelection, SamplingConfig,
    )

    os.makedirs(args.out_dir, exist_ok=True)

    generator = VideoGenerator.from_config(
        GeneratorConfig(
            model_path=args.model_path,
            engine=EngineConfig(
                num_gpus=args.num_gpus,
                use_fsdp_inference=False,
                parallelism=ParallelismConfig(
                    sp_size=1,
                    tp_size=1,
                ),
                offload=OffloadConfig(
                    dit=False,
                    dit_layerwise=False,
                    text_encoder=False,
                    vae=False,
                    image_encoder=False,
                    pin_cpu_memory=False,
                ),
            ),
            pipeline=PipelineSelection(workload_type="t2i"),
        )
    )
    try:
        for i, prompt in enumerate(prompts):
            seed = args.seed + i

            filename_base = f"sd35_{i:02d}_seed{seed}_{_safe_filename(prompt, max_len=80)}"
            _remove_existing_outputs(args.out_dir, filename_base)

            output_path = os.path.join(args.out_dir, f"{filename_base}.png")
            print(f"[sd35] prompt_idx={i} seed={seed} output_path={output_path}")

            generator.generate(
                GenerationRequest(
                    prompt=prompt,
                    negative_prompt=args.negative,
                    sampling=SamplingConfig(
                        height=args.height,
                        width=args.width,
                        num_frames=1,
                        fps=1,
                        num_inference_steps=args.steps,
                        guidance_scale=args.guidance,
                        seed=seed,
                    ),
                    output=OutputConfig(
                        output_path=output_path,
                        save_video=True,
                    ),
                )
            )

        print(f"[sd35] done. outputs written to: {args.out_dir}")
    finally:
        generator.shutdown()


if __name__ == "__main__":
    main()