import numpy as np
from rembg import remove
from PIL import Image, ImageFilter, ImageColor
from moviepy.editor import ImageClip, CompositeVideoClip
import os
import math

class IconAnimator:
    def __init__(self, input_path, output_path):
        self.input_path = input_path
        self.output_path = output_path
        
    def hex_to_rgb(self, hex_color):
        # Converts hex string (e.g., "#00ffc8") to RGB tuple
        try:
            return ImageColor.getcolor(hex_color, "RGB")
        except:
            return (255, 255, 255) # Fallback to white

    def process_image(self, icon_color_hex="#FFFFFF"):
        # 1. Remove Background
        with open(self.input_path, 'rb') as i:
            input_data = i.read()
            output_data = remove(input_data)
        
        # Load into PIL
        from io import BytesIO
        img = Image.open(BytesIO(output_data)).convert("RGBA")
        
        # 2. Recolor the Icon Body
        r, g, b, a = img.split()
        target_r, target_g, target_b = self.hex_to_rgb(icon_color_hex)
        solid_color = Image.new("RGB", img.size, (target_r, target_g, target_b))
        img_colored = Image.merge("RGBA", (*solid_color.split(), a))
        
        return img_colored

    def create_glow_layer(self, img, blur_radius, color_hex):
        # Create padding so glow isn't cut off
        padding = 60
        new_size = (img.width + padding*2, img.height + padding*2)
        canvas = Image.new("RGBA", new_size, (0, 0, 0, 0))
        canvas.paste(img, (padding, padding), img)
        
        # Extract Alpha to create the glow shape
        r, g, b, a = canvas.split()
        
        # Create solid color for glow
        target_r, target_g, target_b = self.hex_to_rgb(color_hex)
        glow_base = Image.new("RGB", new_size, (target_r, target_g, target_b))
        glow_shape = Image.merge("RGBA", (*glow_base.split(), a))
        
        # Apply Blur
        glow = glow_shape.filter(ImageFilter.GaussianBlur(radius=blur_radius))
        return glow, padding

    def generate_animation(self, icon_color="#FFFFFF", glow_color="#00ffc8", speed=2.0, intensity=1.0):
        # 1. Prepare Base Image
        base_pil = self.process_image(icon_color)
        
        # 2. Create Glow Layers (Inner and Outer)
        glow_small, pad = self.create_glow_layer(base_pil, blur_radius=10, color_hex=glow_color) 
        glow_large, _ = self.create_glow_layer(base_pil, blur_radius=30, color_hex=glow_color) 
        
        w, h = base_pil.size
        final_w, final_h = w + (pad*2), h + (pad*2)

        # 3. Create MoviePy Clips
        # Center the base icon
        base_clip = ImageClip(np.array(base_pil)).set_duration(speed).set_position("center")
        glow_s_clip = ImageClip(np.array(glow_small)).set_duration(speed)
        glow_l_clip = ImageClip(np.array(glow_large)).set_duration(speed)

        # 4. Apply Animation (THE FIX)
        # Instead of set_opacity(func), we modify the mask directly using fl()
        
        def pulse_small(t):
            # Oscillates between 0.6 and 1.0
            val = 0.6 + 0.4 * math.sin(2 * math.pi * t / speed)
            return max(0.0, min(val * intensity, 1.0))

        def pulse_large(t):
            # Oscillates between 0.4 and 0.7
            val = 0.4 + 0.3 * math.sin(2 * math.pi * t / speed)
            return max(0.0, min(val * intensity, 1.0))

        # We multiply the existing mask frame (gf(t)) by the pulse value
        glow_s_clip.mask = glow_s_clip.mask.fl(lambda gf, t: gf(t) * pulse_small(t))
        glow_l_clip.mask = glow_l_clip.mask.fl(lambda gf, t: gf(t) * pulse_large(t))

        # 5. Composite
        final = CompositeVideoClip([
            glow_l_clip,   # Outer glow (Background)
            glow_s_clip,   # Inner glow (Rim light)
            base_clip      # Sharp Icon (Foreground)
        ], size=(final_w, final_h))

        # 6. Export
        final.write_gif(self.output_path, fps=15, program='ffmpeg', opt='OptimizeTransparency')
        return self.output_path