import gradio as gr
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import math
import dlib
import tempfile
import requests
import os
from transformers import pipeline
import cv2
import io
import json
import re
import time

detector = dlib.get_frontal_face_detector()
try:
    predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
except RuntimeError:
    print("Downloading shape_predictor_68_face_landmarks.dat...")
    landmarks_url = "http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2"
    landmarks_compressed = requests.get(landmarks_url).content
    import bz2
    landmarks_data = bz2.decompress(landmarks_compressed)
    with open("shape_predictor_68_face_landmarks.dat", "wb") as f:
        f.write(landmarks_data)
    predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")

IMAGE_GEN_API = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
HF_TOKEN = os.getenv("HF_TOKEN")

LLM_API = "https://api-inference.huggingface.co/models/lmsys/fastchat-t5-3b-v1.0"

def query_hf_image_generation(prompt):
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    payload = {"inputs": prompt}
    for _ in range(3):
        try:
            response = requests.post(IMAGE_GEN_API, headers=headers, json=payload)
            response.raise_for_status()
            image_bytes = response.content
            image = Image.open(io.BytesIO(image_bytes))
            return image
        except requests.exceptions.RequestException as e:
            print(f"Image generation attempt failed: {e}")
            time.sleep(2)
    raise Exception("Image generation failed after multiple attempts")

def query_llm(prompt, system_prompt):
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    prompt_template = f"<|system|>\n{system_prompt}</s>\n<|user|>\n{prompt}</s>\n<|assistant|>\n"
    payload = {"inputs": prompt_template, "max_new_tokens": 200}
    for _ in range(3):
        try:
            response = requests.post(LLM_API, headers=headers, json=payload)
            response.raise_for_status()
            return response.json()[0]['generated_text']
        except requests.exceptions.RequestException as e:
            print(f"LLM query attempt failed: {e}")
            time.sleep(2)
    raise Exception("LLM query failed after multiple attempts")

def segment_script(script):
    system_prompt = "You are a helpful assistant. Given a script, divide it into segments suitable for generating images, ensuring each segment is less than 500 characters."
    llm_response = query_llm(script, system_prompt)
    segments = llm_response.split('\n')
    segments = [seg.strip() for seg in segments if seg.strip()]
    return segments

def generate_image_prompts(script_segments):
    image_prompts = []
    for segment in script_segments:
        system_prompt = "You are a helpful assistant. Create a concise image prompt based on the following script segment:"
        prompt = f"Script Segment: {segment}"
        image_prompt = query_llm(prompt, system_prompt)
        image_prompts.append(image_prompt)
    return image_prompts

def extract_motion_params(llm_output):
    try:
        start_index = llm_output.find('{')
        end_index = llm_output.rfind('}') + 1
        json_string = llm_output[start_index:end_index]
        params = json.loads(json_string)
        return params
    except:
        return {
            "motion_type": "none",
            "intensity": 0.25,
            "text_overlay": "",
            "text_color": "white",
            "start_time": 0,
            "end_time": 5
        }

def detect_face_landmarks(image):
    gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
    rects = detector(gray, 1)
    if len(rects) > 0:
        shape = predictor(gray, rects[0])
        shape = np.array([(shape.part(i).x, shape.part(i).y) for i in range(68)])
        return shape
    else:
        return None

def apply_color_grading(frame, color_preset, intensity):
    if color_preset == "sepia":
        sepia_matrix = np.array([[0.393, 0.769, 0.189],
                                [0.349, 0.686, 0.168],
                                [0.272, 0.534, 0.131]])
        frame_float = frame.astype(np.float32) / 255.0
        sepia_effect = cv2.transform(frame_float, sepia_matrix)
        blended_frame = (1 - intensity) * frame_float + intensity * sepia_effect
        return (np.clip(blended_frame, 0, 1) * 255).astype(np.uint8)
    elif color_preset == "vintage":
        frame_float = frame.astype(np.float32) / 255.0
        frame_float[:, :, 0] *= (1 - intensity * 0.6)
        frame_float[:, :, 2] *= (1 + intensity * 0.3)
        grayscale = cv2.cvtColor(frame_float, cv2.COLOR_RGB2GRAY)
        grayscale_rgb = cv2.cvtColor(grayscale, cv2.COLOR_GRAY2RGB)
        blended_frame = (1 - intensity * 0.5) * frame_float + intensity * 0.5 * grayscale_rgb
        return (np.clip(blended_frame, 0, 1) * 255).astype(np.uint8)
    elif color_preset == "black_and_white":
        gray_frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
        return cv2.cvtColor(gray_frame, cv2.COLOR_GRAY2RGB)
    elif color_preset == "cold":
        frame_float = frame.astype(np.float32) / 255.0
        frame_float[:, :, 0] *= (1 + intensity * 0.7)
        frame_float[:, :, 2] *= (1 - intensity * 0.2)
        return (np.clip(frame_float, 0, 1) * 255).astype(np.uint8)
    elif color_preset == "warm":
        frame_float = frame.astype(np.float32) / 255.0
        frame_float[:, :, 2] *= (1 + intensity * 0.7)
        frame_float[:, :, 0] *= (1 - intensity * 0.2)
        return (np.clip(frame_float, 0, 1) * 255).astype(np.uint8)
    elif color_preset == "neon":
        frame_float = frame.astype(np.float32) / 255.0
        lab = cv2.cvtColor(frame_float, cv2.COLOR_RGB2LAB)
        l, a, b = cv2.split(lab)
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
        l = clahe.apply(l)
        lab = cv2.merge((l, a, b))
        frame_float = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
        frame_float[:, :, 0] *= (1 - intensity * 0.4) 
        frame_float[:, :, 1] *= (1 + intensity * 0.8) 
        frame_float[:, :, 2] *= (1 - intensity * 0.4)
        return (np.clip(frame_float, 0, 1) * 255).astype(np.uint8)

    return frame

def apply_vignette(frame, intensity):
    width, height = frame.shape[1], frame.shape[0]
    x = np.linspace(-1, 1, width)
    y = np.linspace(-1, 1, height)
    X, Y = np.meshgrid(x, y)
    radius = np.sqrt(X**2 + Y**2)
    vignette = 1 - intensity * radius**2
    vignette = np.clip(vignette, 0, 1)
    vignette = np.stack([vignette] * 3, axis=-1)
    frame_float = frame.astype(np.float32) / 255.0
    result = frame_float * vignette
    return (np.clip(result, 0, 1) * 255).astype(np.uint8)

def apply_bokeh(frame, intensity, t):
    frame_float = frame.astype(np.float32) / 255.0
    gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
    circles = []
    for _ in range(int(intensity * 30)):
        radius = np.random.randint(5, 30)
        x = np.random.randint(radius, frame.shape[1] - radius)
        y = np.random.randint(radius, frame.shape[0] - radius)
        color = frame_float[y, x]
        brightness = np.random.uniform(0.5, 1.0)
        circles.append((x, y, radius, color, brightness))
    
    bokeh_effect = np.zeros_like(frame_float)
    for x, y, radius, color, brightness in circles:
        y_grid, x_grid = np.ogrid[-y:frame.shape[0]-y, -x:frame.shape[1]-x]
        mask = x_grid*x_grid + y_grid*y_grid <= radius*radius
        bokeh_effect[mask] += np.array(color) * brightness * (0.5 + 0.5 * np.sin(t * 2 * math.pi))

    blended_frame = frame_float + intensity * bokeh_effect
    return (np.clip(blended_frame, 0, 1) * 255).astype(np.uint8)

def apply_advanced_motion(image, motion_type, intensity, duration, fps, text_overlay, text_color, font_size, start_time, end_time, color_preset, vignette_intensity):
    frames = []
    width, height = image.size
    landmarks = detect_face_landmarks(image)

    for i in range(int(duration * fps)):
        t = i / (duration * fps)
        frame = image.copy()

        if landmarks is not None:
            if motion_type == "head_nod":
                top_head = landmarks[27]
                bottom_head = landmarks[8]
                angle = math.sin(t * 2 * math.pi) * intensity * 8
                center_x = (top_head[0] + bottom_head[0]) // 2
                center_y = (top_head[1] + bottom_head[1]) // 2
                M = cv2.getRotationMatrix2D((center_x, center_y), angle, 1)
                rotated_image = cv2.warpAffine(np.array(image), M, (width, height), flags=cv2.INTER_LANCZOS4)
                frame = Image.fromarray(rotated_image)

            elif motion_type == "head_shake":
                top_head = landmarks[27]
                left_head = landmarks[0]
                right_head = landmarks[16]
                angle = math.sin(t * 3 * math.pi) * intensity * 6
                center_x = top_head[0]
                center_y = top_head[1]
                M = cv2.getRotationMatrix2D((center_x, center_y), angle, 1)
                rotated_image = cv2.warpAffine(np.array(image), M, (width, height), flags=cv2.INTER_LANCZOS4)
                frame = Image.fromarray(rotated_image)

            elif motion_type == "eye_blink":
                left_eye_top = landmarks[37]
                left_eye_bottom = landmarks[41]
                right_eye_top = landmarks[43]
                right_eye_bottom = landmarks[47]
                blink_progress = abs(math.sin(t * 2 * math.pi))
                if blink_progress > 0.9:
                    draw = ImageDraw.Draw(frame)
                    draw.line([tuple(landmarks[36]), tuple(landmarks[39])], fill=text_color, width=2)
                    draw.line([tuple(landmarks[42]), tuple(landmarks[45])], fill=text_color, width=2)
                else:
                    frame = image.copy()

            elif motion_type == "smile":
                mouth_left = landmarks[48]
                mouth_right = landmarks[54]
                mouth_top = landmarks[51]
                mouth_bottom = landmarks[57]
                smile_progress = intensity * t
                
                draw = ImageDraw.Draw(frame)
                curve_points = [
                    tuple(mouth_left),
                    (mouth_left[0] + (mouth_right[0] - mouth_left[0]) // 4, mouth_left[1] + int(20 * smile_progress)),
                    (mouth_left[0] + 3 * (mouth_right[0] - mouth_left[0]) // 4, mouth_right[1] + int(20 * smile_progress)),
                    tuple(mouth_right)
                ]
                draw.line(curve_points, fill=text_color, width=4)

        if motion_type == "zoom":
            scale = 1 + intensity * t
            new_size = (int(width * scale), int(height * scale))
            resized_image = image.resize(new_size, Image.Resampling.LANCZOS)
            x_offset = (new_size[0] - width) // 2
            y_offset = (new_size[1] - height) // 2
            frame = resized_image.crop((x_offset, y_offset, x_offset + width, y_offset + height))

        elif motion_type == "pan":
            x_offset = int(intensity * t * (width - width))
            y_offset = int(intensity * t * (height - height))
            frame = Image.new("RGB", (width, height))
            frame.paste(image, (-x_offset, -y_offset))

        elif motion_type == "rotate":
            angle = intensity * t * 360
            rotated_image = image.rotate(angle, expand=True, resample=Image.Resampling.BICUBIC)
            x_offset = (rotated_image.width - width) // 2
            y_offset = (rotated_image.height - height) // 2
            frame = Image.new("RGB", (width, height))
            frame.paste(rotated_image, (-x_offset, -y_offset))

        elif motion_type == "move_right":
            x_offset = int(intensity * t * width)
            frame = Image.new("RGB", (width, height), "black")
            frame.paste(image, (x_offset, 0))

        elif motion_type == "move_left":
            x_offset = -int(intensity * t * width)
            frame = Image.new("RGB", (width, height), "black")
            frame.paste(image, (x_offset, 0))

        elif motion_type == "move_up":
            y_offset = -int(intensity * t * height)
            frame = Image.new("RGB", (width, height), "black")
            frame.paste(image, (0, y_offset))

        elif motion_type == "move_down":
            y_offset = int(intensity * t * height)
            frame = Image.new("RGB", (width, height), "black")
            frame.paste(image, (0, y_offset))
        
        elif motion_type == "shake":
            shake_intensity = intensity * 10  
            x_offset = int(shake_intensity * math.sin(t * 2 * math.pi * 5))  
            y_offset = int(shake_intensity * math.cos(t * 2 * math.pi * 3))  
            frame = Image.new("RGB", (width, height))
            frame.paste(image, (x_offset, y_offset))
        
        elif motion_type == "fade_in":
            alpha = t 
            frame = Image.blend(Image.new("RGB", (width, height), "black"), image, alpha)

        elif motion_type == "fade_out":
            alpha = 1 - t
            frame = Image.blend(Image.new("RGB", (width, height), "black"), image, alpha)
        
        elif motion_type == "rain":
            draw = ImageDraw.Draw(frame)
            for _ in range(int(intensity * 5)):
                x = np.random.randint(0, width)
                y = np.random.randint(0, height)
                length = np.random.randint(5, 15)
                speed = intensity * 3
                y_end = y + length + i * speed
                draw.line([(x, y), (x, y_end)], fill="lightblue", width=1)
        
        elif motion_type == "bokeh":
            frame_np = np.array(frame)
            frame_np = apply_bokeh(frame_np, intensity, t)
            frame = Image.fromarray(frame_np)

        frame_np = np.array(frame)
        
        if color_preset:
            frame_np = apply_color_grading(frame_np, color_preset, intensity)
        if vignette_intensity > 0:
            frame_np = apply_vignette(frame_np, vignette_intensity)

        frame = Image.fromarray(frame_np)

        draw = ImageDraw.Draw(frame)
        if text_overlay and start_time <= t <= end_time:
            try:
                font = ImageFont.truetype("arial.ttf", font_size)
            except IOError:
                font = ImageFont.load_default()
            text_width, text_height = draw.textsize(text_overlay, font=font)
            x = (width - text_width) // 2
            y = (height - text_height) // 2
            draw.text((x, y), text_overlay, font=font, fill=text_color)

        frames.append(np.array(frame))

    return frames

def create_video_from_frames(frames, output_filename, fps=30):
    writer = imageio_ffmpeg.write_frames(output_filename, frames[0].shape[:2], pix_fmt_out='yuv420p', fps=fps, codec='libx264', preset="veryslow")
    writer.send(None)
    for frame in frames:
        writer.send(frame)
    writer.close()

def generate_video_from_script(script, duration_per_segment=5):
    script_segments = segment_script(script)
    image_prompts = generate_image_prompts(script_segments)
    all_frames = []

    for i, (segment, image_prompt) in enumerate(zip(script_segments, image_prompts)):
        print(f"Processing segment {i + 1} of {len(script_segments)}")
        print(f"  Segment: {segment}")
        print(f"  Image Prompt: {image_prompt}")
        
        image = query_hf_image_generation(image_prompt)
        image_description = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")(image)[0]['generated_text']
        
        system_prompt = "You are an expert in image to video creation. Provide the motion type, intensity, text overlay, text color, text start and end times, color preset, and vignette intensity for the following image description and user prompt. Give the response in a JSON format."
        prompt = f"Image Description: {image_description}\nUser Prompt: {segment}"
        llm_response = query_llm(prompt, system_prompt)
        
        print(f"  LLM Response: {llm_response}")
        motion_params = extract_motion_params(llm_response)
        print(f"  Motion Parameters: {motion_params}")
        
        frames = apply_advanced_motion(
            image,
            motion_params["motion_type"],
            motion_params["intensity"],
            duration=duration_per_segment,
            fps=30,
            text_overlay=motion_params["text_overlay"],
            text_color=motion_params["text_color"],
            font_size=50,
            start_time=motion_params["start_time"],
            end_time=motion_params["end_time"],
            color_preset=motion_params.get("color_preset", None),
            vignette_intensity=motion_params.get("vignette_intensity", 0)
        )
        all_frames.extend(frames)

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
        output_filename = tmpfile.name
    create_video_from_frames(all_frames, output_filename)
    return output_filename

def generate_and_animate(prompt):
    try:
        image = query_hf_image_generation(prompt)
        image_description = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")(image)[0]['generated_text']
        llm_response = query_llm(prompt, image_description)
        motion_params = extract_motion_params(llm_response)
        frames = apply_advanced_motion(
            image,
            motion_params["motion_type"],
            motion_params["intensity"],
            duration=5,
            fps=30,
            text_overlay=motion_params["text_overlay"],
            text_color=motion_params["text_color"],
            font_size=50,
            start_time=motion_params["start_time"],
            end_time=motion_params["end_time"],
            color_preset=motion_params.get("color_preset", None),
            vignette_intensity=motion_params.get("vignette_intensity", 0)
        )
        video_file = create_video_from_frames(frames)
        return video_file, gr.Image.update(value=image)
    except Exception as e:
        return str(e), None

motion_types = [
    "zoom", "pan", "rotate", "move_right", "move_left", "move_up", "move_down", 
    "shake", "fade_in", "fade_out", "head_nod", "head_shake", "eye_blink", "smile", "rain", "bokeh", "none"
]
text_colors = ["white", "black", "red", "green", "blue", "yellow"]
color_presets = ["sepia", "vintage", "black_and_white", "cold", "warm", "neon", "none"]

iface = gr.Interface(
    fn=generate_and_animate,
    inputs=[
        gr.Textbox(label="Prompt"),
    ],
    outputs=[
        gr.Video(label="Generated Video"),
        gr.Image(label="Generated Image")
    ],
    title="AI Video Generator",
    description="Enter a prompt to generate an image and animate it. Uses Flux 1, an LLM, and advanced video processing techniques."
)

video_iface = gr.Interface(
    fn=generate_video_from_script,
    inputs=[
        gr.Textbox(label="Script (max 1 minute video)", lines=5),
        gr.Slider(label="Duration per Segment (seconds)", minimum=1, maximum=10, step=1, value=5)
    ],
    outputs=gr.Video(label="Generated Video from Script"),
    title="Story Visualizer",
    description="Enter a short story script, and this will generate a video visualizing it using multiple images and animations."
)
demo = gr.TabbedInterface([iface, video_iface], ["Generate and Animate", "Story to Video"])

if __name__ == "__main__":
    demo.launch(share=True, debug=True)