from __future__ import annotations

from pathlib import Path
import shutil

from model.Animation.animation_loader import (
    build_animation_predictor,
    UnifiedAnimationConfig,
    AnimationBackend,
)


# -----------------------------
# 配置区（按需修改）
# -----------------------------
OUTPUT_DIR = "outputs/test_animation"
ANIMATION_BACKEND = AnimationBackend.ANIMATEDIFF
OUTPUT_FORMAT = "png_sequence"  # "gif" | "png_sequence"

PROMPT = "a cinematic mountain landscape, camera slowly pans left"
NEGATIVE_PROMPT = "blurry, low quality"
NUM_INFERENCE_STEPS = 25
GUIDANCE_SCALE = 8.0
WIDTH = 512
HEIGHT = 512
VIDEO_LENGTH = 16
SEED = -1
CONTROL_IMAGE_PATH = "path/to/your_image.png"


def main() -> None:
    base_dir = Path(__file__).resolve().parent
    out_dir = base_dir / OUTPUT_DIR
    out_dir.mkdir(parents=True, exist_ok=True)

    predictor, used_backend = build_animation_predictor(
        UnifiedAnimationConfig(backend=ANIMATION_BACKEND)
    )

    if CONTROL_IMAGE_PATH.strip() in {"", "path/to/your_image.png"}:
        raise ValueError("请先设置 CONTROL_IMAGE_PATH 为你的输入图片路径（png/jpg）。")

    control_image = (base_dir / CONTROL_IMAGE_PATH).resolve()
    if not control_image.is_file():
        raise FileNotFoundError(f"control image not found: {control_image}")

    result_path = predictor(
        prompt=PROMPT,
        negative_prompt=NEGATIVE_PROMPT,
        num_inference_steps=NUM_INFERENCE_STEPS,
        guidance_scale=GUIDANCE_SCALE,
        width=WIDTH,
        height=HEIGHT,
        video_length=VIDEO_LENGTH,
        seed=SEED,
        control_image_path=str(control_image),
        output_format=OUTPUT_FORMAT,
    )

    source = Path(result_path)
    if OUTPUT_FORMAT == "png_sequence":
        out_seq_dir = out_dir / f"{used_backend.value}_frames"
        if out_seq_dir.exists():
            shutil.rmtree(out_seq_dir)
        shutil.copytree(source, out_seq_dir)
        print(f"[Animation] backend={used_backend.value}, saved={out_seq_dir}")
        return

    out_path = out_dir / f"{used_backend.value}.gif"
    out_path.write_bytes(source.read_bytes())
    print(f"[Animation] backend={used_backend.value}, saved={out_path}")


if __name__ == "__main__":
    main()