#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from __future__ import annotations
import os
from pathlib import Path
from typing import List, Optional

import torch
from PIL import Image

from diffusers.models import (
    AutoencoderKL,
    ControlNetModel,
)

from diffusers import (
    DPMSolverMultistepScheduler,
    DDIMScheduler,
)

from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.models.unets.unet_motion_model import MotionAdapter




APP_HOME = Path(os.environ.get("APP_HOME", "/app"))
OUT_DIR = APP_HOME / "outputs" / "anime_service"
OUT_DIR.mkdir(parents=True, exist_ok=True)

# IDs padrão (editáveis na UI)
MOTION_ID = os.environ.get("MOTION_ID", "guoyww/animatediff-motion-adapter-v1-5-2")
CONTROLNET_ID = os.environ.get("CONTROLNET_ID", "lllyasviel/control_v11p_sd15_openpose")
VAE_ID = os.environ.get("VAE_ID", "stabilityai/sd-vae-ft-mse")
MODEL_ID = os.environ.get("MODEL_ID", "SG161222/Realistic_Vision_V5.1_noVAE")

TORCH_DTYPE = (
    torch.float16
    if os.environ.get("TORCH_DTYPE", "bfloat16").lower() in ("fp16", "float16", "half")
    else torch.bfloat16
)

class AnimateDiffControlNetService:
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.pipe: DiffusionPipeline | None = None

    def setup(
        self,
        model_id: str = MODEL_ID,
        motion_id: str = MOTION_ID,
        controlnet_id: str = CONTROLNET_ID,
        vae_id: str = VAE_ID,
        dtype: torch.dtype = TORCH_DTYPE,
    ) -> str:
        # Carrega adaptador de movimento, controlnet e VAE
        adapter = MotionAdapter.from_pretrained(motion_id)
        controlnet = ControlNetModel.from_pretrained(controlnet_id, torch_dtype=dtype)
        vae = AutoencoderKL.from_pretrained(vae_id, torch_dtype=dtype)

        # Monta pipeline custom AnimateDiff + ControlNet
        pipe = DiffusionPipeline.from_pretrained(
            model_id,
            motion_adapter=adapter,
            controlnet=controlnet,
            vae=vae,
            custom_pipeline="pipeline_animatediff_controlnet",  # requer arquivo/pipeline correspondente
        ).to(device=self.device, dtype=dtype)

        # Scheduler robusto: tenta DPMSolverMultistep; se config estiver incompatível, cai para DDIM
        try:
            pipe.scheduler = DPMSolverMultistepScheduler.from_pretrained(
                model_id,
                subfolder="scheduler",
                clip_sample=False,
                timestep_spacing="linspace",
                steps_offset=1,
                beta_schedule="linear",
            )
        except Exception:
            pipe.scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")

        # Slicing de VAE
        try:
            pipe.vae.enable_slicing()
        except Exception:
            pass

        self.pipe = pipe
        return f"AnimeDiff pronto em {self.device} (dtype={dtype})"

    @torch.inference_mode()
    def generate_frames(
        self,
        conditioning_frames: List[Image.Image],
        prompt: str,
        negative_prompt: str = "",
        width: int = 512,
        height: int = 768,
        num_inference_steps: int = 12,
        guidance_scale: float = 7.5,
        seed: Optional[int] = None,
    ) -> List[str]:
        assert self.pipe is not None, "Chame setup() antes."

        gen = None
        if seed is not None and seed >= 0:
            gen = torch.Generator(device=self.device).manual_seed(int(seed))

        result = self.pipe(
            prompt=prompt,
            negative_prompt=(negative_prompt or None),
            width=int(width),
            height=int(height),
            conditioning_frames=conditioning_frames,
            num_inference_steps=int(num_inference_steps),
            guidance_scale=float(guidance_scale),
            generator=gen,
            output_type="pil",
        )

        frames = result.frames[0]
        out_dir = OUT_DIR / "session"
        out_dir.mkdir(parents=True, exist_ok=True)
        paths: List[str] = []
        for i, img in enumerate(frames):
            p = out_dir / f"frame_{i:02d}.png"
            img.save(p)
            paths.append(str(p))
        return paths