import torch
import numpy as np
from diffusers import AutoPipelineForImage2Image, LCMScheduler, EulerAncestralDiscreteScheduler
from PIL import Image
import utils
import os
import uuid
import zipfile
import cv2
import gc

class DeforumRunner:
    def __init__(self, device="cpu"):
        self.device = device
        self.pipe = None
        self.stop_requested = False
        self.current_model_config = (None, None, None)

    def load_model(self, model_id, lora_id, scheduler_name):
        """Loads model, LoRA, and scheduler dynamically."""
        new_config = (model_id, lora_id, scheduler_name)
        if new_config == self.current_model_config and self.pipe is not None:
            return

        print(f"Loading Model: {model_id}, LoRA: {lora_id}, Scheduler: {scheduler_name}")
        if self.pipe:
            del self.pipe
            gc.collect()

        try:
            pipe = AutoPipelineForImage2Image.from_pretrained(
                model_id, safety_checker=None, torch_dtype=torch.float32
            )
        except Exception as e:
            print(f"Error loading model {model_id}: {e}. Falling back.")
            pipe = AutoPipelineForImage2Image.from_pretrained(
                "runwayml/stable-diffusion-v1-5", safety_checker=None, torch_dtype=torch.float32
            )

        if lora_id and lora_id != "None":
            try:
                pipe.load_lora_weights(lora_id)
                pipe.fuse_lora()
                print("LoRA loaded and fused.")
            except Exception as e: print(f"Error loading LoRA: {e}")

        if scheduler_name == "LCM":
            pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
        elif scheduler_name == "Euler A":
            pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)

        pipe.to(self.device)
        pipe.set_progress_bar_config(disable=True)
        pipe.enable_attention_slicing() # Crucial for CPU memory
        self.pipe = pipe
        self.current_model_config = new_config
        print("Pipeline ready.")

    def stop(self):
        self.stop_requested = True

    def render(self, prompts, neg_prompt, max_frames, width, height,
               zoom_s, angle_s, tx_s, ty_s, strength_s, noise_s, 
               fps, steps, cadence, color_mode, border_mode, init_image_upload,
               model_id, lora_id, scheduler_name):
        
        self.stop_requested = False
        self.load_model(model_id, lora_id, scheduler_name)
        
        # 1. Parse Schedules
        keys = ['z', 'a', 'tx', 'ty', 'str', 'noi']
        inputs = [zoom_s, angle_s, tx_s, ty_s, strength_s, noise_s]
        sched = {k: utils.parse_weight_string(v, max_frames) for k, v in zip(keys, inputs)}
        
        # 2. Setup
        run_id = uuid.uuid4().hex[:6]
        output_dir = f"output_{run_id}"
        os.makedirs(output_dir, exist_ok=True)
        
        if init_image_upload:
            prev_img = init_image_upload.resize((width, height), Image.LANCZOS)
            color_anchor = prev_img
        else:
            prev_img = None
            color_anchor = None
            
        generated_frames = []
        print(f"Starting run {run_id}...")

        # 3. Loop
        for i in range(max_frames):
            if self.stop_requested:
                print("Generation stopped.")
                break
            
            # Get Params
            z, a, tx, ty = sched['z'][i], sched['a'][i], sched['tx'][i], sched['ty'][i]
            strength, noise = sched['str'][i], sched['noi'][i]
            current_prompt = prompts[max(k for k in prompts.keys() if k <= i)]
            
            # --- Authentic Deforum Loop ---
            # 1. Warp Previous Frame
            if prev_img is not None:
                warped_img = utils.anim_frame_warp_2d(prev_img, {'angle': a, 'zoom': z, 'translation_x': tx, 'translation_y': ty}, border_mode)
            else:
                warped_img = Image.new("RGB", (width, height), (0,0,0))

            # Decide: Diffusion or Just Warp (Cadence)
            if i % cadence == 0:
                # 2. Color Match & 3. Add Noise (Only before diffusion)
                init_for_diff = utils.maintain_colors(warped_img, color_anchor, color_mode)
                init_for_diff = utils.add_noise(init_for_diff, noise)
                
                # 4. Diffusion (Img2Img)
                # Use high strength for first frame if no init image provided
                curr_strength = strength if prev_img is not None else 0.95
                
                gen_image = self.pipe(
                    prompt=current_prompt, negative_prompt=neg_prompt,
                    image=init_for_diff, num_inference_steps=steps,
                    strength=curr_strength, guidance_scale=1.2, # Low CFG for LCM
                    width=width, height=height
                ).images[0]
            else:
                # Cadence step: Just show the warped image (faster)
                gen_image = warped_img

            # Update state
            prev_img = gen_image
            if color_anchor is None: color_anchor = gen_image
            generated_frames.append(gen_image)
            yield gen_image, None, None

        # 4. Finalize
        vid_path = f"{output_dir}/video.mp4"
        self.save_video(generated_frames, vid_path, fps)
        zip_path = f"{output_dir}/frames.zip"
        self.save_zip(generated_frames, zip_path)
        yield generated_frames[-1], vid_path, zip_path

    # (save_video and save_zip are the same as before, omitted for brevity)
    def save_video(self, frames, path, fps):
        if not frames: return
        w, h = frames[0].size
        out = cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
        for f in frames: out.write(cv2.cvtColor(np.array(f), cv2.COLOR_RGB2BGR))
        out.release()
    def save_zip(self, frames, path):
        import io
        with zipfile.ZipFile(path, 'w') as zf:
            for i, f in enumerate(frames):
                buf = io.BytesIO()
                f.save(buf, format="PNG")
                zf.writestr(f"{i:05d}.png", buf.getvalue())