from __future__ import annotations

import datetime
import shutil
from typing import Any, Dict

from fastapi import APIRouter
from PIL import Image

from ..deps import OUTPUT_DIR
from ..schemas import AnimateRequest, CharacterAnimateRequest
from ..services.animation_frames import (
    align_size_to_input,
    composite_rgba_on_background,
    list_png_frames,
    parse_background_color,
    remove_background_by_color,
)
from ..services.image_io import b64_to_pil_image, pil_image_to_png_b64
from ..services.predictors import get_animation_predictor

router = APIRouter(tags=["animation"])


@router.post("/animate")
def animate(req: AnimateRequest) -> Dict[str, Any]:
    try:
        model_name = req.model_name or "animatediff"
        predictor = get_animation_predictor(model_name)

        result_path = predictor(
            prompt=req.prompt,
            negative_prompt=req.negative_prompt or "",
            num_inference_steps=req.num_inference_steps,
            guidance_scale=req.guidance_scale,
            width=req.width,
            height=req.height,
            video_length=req.video_length,
            seed=req.seed,
        )

        out_dir = OUTPUT_DIR / "animation"
        out_dir.mkdir(parents=True, exist_ok=True)
        ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        out_path = out_dir / f"{model_name}_{ts}.gif"
        shutil.copy2(result_path, out_path)

        return {"success": True, "output_path": str(out_path)}
    except Exception as e:
        return {"success": False, "error": str(e)}


@router.post("/animate/character_sequence")
def animate_character_sequence(req: CharacterAnimateRequest) -> Dict[str, Any]:
    """
    输入透明背景角色 PNG，生成 30fps 角色动画 PNG 序列。

    流程：
    - 将角色合成到指定纯色背景，作为 AnimateDiff 控制图
    - 以与输入相近、且对齐到 8 的分辨率生成 PNG 序列
    - 对每帧按背景色做色键剔除，返回透明背景 PNG 序列给前端
    """
    try:
        model_name = req.model_name or "animatediff"
        background_rgb = parse_background_color(req.background_color)

        character = b64_to_pil_image(req.image_b64).convert("RGBA")
        orig_w, orig_h = character.size
        run_w, run_h = align_size_to_input(orig_w, orig_h, req.max_side)

        predictor = get_animation_predictor(model_name)

        out_dir = OUTPUT_DIR / "animation" / "character_sequence"
        out_dir.mkdir(parents=True, exist_ok=True)
        ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        seq_dir = out_dir / f"{model_name}_{ts}"
        seq_dir.mkdir(parents=True, exist_ok=True)

        control_path = seq_dir / "control_image.png"
        control_image = composite_rgba_on_background(
            character,
            background_rgb,
            size=(run_w, run_h),
        )
        control_image.save(control_path)

        frame_dir = predictor(
            prompt=req.prompt,
            negative_prompt=req.negative_prompt or "",
            num_inference_steps=req.num_inference_steps,
            guidance_scale=req.guidance_scale,
            width=run_w,
            height=run_h,
            video_length=req.video_length,
            seed=req.seed,
            control_image_path=str(control_path),
            output_format="png_sequence",
        )

        frame_paths = list_png_frames(frame_dir)
        if not frame_paths:
            raise FileNotFoundError(f"未找到 AnimateDiff 输出 PNG 序列: {frame_dir}")

        frames: list[dict[str, Any]] = []
        processed_dir = seq_dir / "frames"
        processed_dir.mkdir(parents=True, exist_ok=True)
        for idx, frame_path in enumerate(frame_paths):
            frame = remove_background_by_color(
                Image.open(frame_path),
                background_rgb,
                tolerance=req.background_tolerance,
                output_size=(orig_w, orig_h),
            )
            out_path = processed_dir / f"{idx:04d}.png"
            frame.save(out_path)
            frames.append(
                {
                    "index": idx,
                    "image_b64": pil_image_to_png_b64(frame),
                    "path": str(out_path),
                }
            )

        return {
            "success": True,
            "fps": 30,
            "frame_count": len(frames),
            "width": orig_w,
            "height": orig_h,
            "background_color": req.background_color,
            "output_dir": str(processed_dir),
            "frames": frames,
        }
    except Exception as e:
        return {"success": False, "error": str(e), "frames": []}