# Import necessary libraries
from kokoro import KPipeline
import soundfile as sf
import os
from moviepy.editor import (
    VideoFileClip, concatenate_videoclips, AudioFileClip, ImageClip,
    CompositeVideoClip, TextClip
)
import moviepy.video.fx.all as vfx
import moviepy.config as mpy_config
from pydub import AudioSegment
from PIL import Image
import tempfile
import random
import cv2
import math
import requests
from gtts import gTTS
import gradio as gr
import numpy as np
import shutil

# Initialize Kokoro TTS pipeline
pipeline = KPipeline(lang_code='a')
# Ensure ImageMagick binary is set (adjust path as needed)
mpy_config.change_settings({"IMAGEMAGICK_BINARY": "/usr/bin/convert"})

# Global Configuration
PEXELS_API_KEY = 'BhJqbcdm9Vi90KqzXKAhnEHGsuFNv4irXuOjWtT761U49lRzo03qBGna'
OPENROUTER_API_KEY = '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'
OPENROUTER_MODEL = "mistralai/mistral-small-3.1-24b-instruct:free"
OUTPUT_VIDEO_FILENAME = "final_video.mp4"
USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
MAX_CLIPS = 10  # Maximum number of editable clips

# Temporary folder setup
TEMP_FOLDER = tempfile.mkdtemp()

# Existing Helper Functions (unchanged)
def generate_script(user_input):
    headers = {
        'Authorization': f'Bearer {OPENROUTER_API_KEY}',
        'HTTP-Referer': 'https://your-domain.com',
        'X-Title': 'AI Documentary Maker'
    }
    prompt = f"""Short Documentary Script GeneratorInstructions:
If I say "use this," just output the script exactly as I gave it.
If I only give topics, generate a script based on them.
If I provide a full script, rewrite it without any changes. Make everything short simple and humarous funny and act as serious but humarous. And don't say anything off topic. Also alway say a funny statement to subscribe based on the video topic at the end. Use normal conversational text like a normal person talking and avoid AI phase make the statements humanize and normal conversational
And always talk negative views don't really wants any positive statement and the full script will be one video about one topic so keep in mind
Formatting Rules:
Title in Square Brackets:
Each section starts with a one-word title inside [ ] (max two words if necessary).
This title will be used as a search term for Pexels footage.
Casual & Funny Narration:
Each section has 5-10 words of narration.
Keep it natural, funny, and unpredictable (not robotic, poetic, or rhythmic).
No Special Formatting:
No bold, italics, or special characters. You are a assistant AI your task is to create script. You aren't a chatbot. So, don't write extra text
Generalized Search Terms:
If a term is too specific, make it more general for Pexels search.
Scene-Specific Writing:
Each section describes only what should be shown in the video.
Output Only the Script, and also make it funny and humarous and helirous and also add to subscribe with a funny statement like subscribe now or .....
No extra text, just the script.
Example Output:
[North Korea]
Top 5 unknown facts about North Korea.
[Invisibility]
North Korea’s internet speed is so fast… it doesn’t exist.
[Leadership]
Kim Jong-un once won an election with 100% votes… against himself.
[Magic]
North Korea discovered time travel. That’s why their news is always from the past.
[Warning]
Subscribe now, or Kim Jong-un will send you a free one-way ticket… to North Korea.
[Freedom]
North Korean citizens can do anything… as long as it's government-approved.
Now here is the Topic/scrip: {user_input}
"""
    data = {
        'model': OPENROUTER_MODEL,
        'messages': [{'role': 'user', 'content': prompt}],
        'temperature': 0.4,
        'max_tokens': 5000
    }
    try:
        response = requests.post(
            'https://openrouter.ai/api/v1/chat/completions',
            headers=headers,
            json=data,
            timeout=30
        )
        if response.status_code == 200:
            return response.json()['choices'][0]['message']['content']
        else:
            print(f"API Error {response.status_code}: {response.text}")
            return None
    except Exception as e:
        print(f"Request failed: {str(e)}")
        return None

def parse_script(script_text):
    sections = {}
    current_title = None
    current_text = ""
    for line in script_text.splitlines():
        line = line.strip()
        if line.startswith("[") and "]" in line:
            bracket_start = line.find("[")
            bracket_end = line.find("]", bracket_start)
            if bracket_start != -1 and bracket_end != -1:
                if current_title is not None:
                    sections[current_title] = current_text.strip()
                current_title = line[bracket_start+1:bracket_end]
                current_text = line[bracket_end+1:].strip()
        elif current_title:
            current_text += line + " "
    if current_title:
        sections[current_title] = current_text.strip()
    clips = [{"title": title, "narration": narration} for title, narration in sections.items()]
    return clips

def search_pexels_videos(query, pexels_api_key):
    headers = {'Authorization': pexels_api_key}
    url = "https://api.pexels.com/videos/search"
    params = {"query": query, "per_page": 15}
    try:
        response = requests.get(url, headers=headers, params=params, timeout=10)
        if response.status_code == 200:
            videos = response.json().get("videos", [])
            hd_videos = [v["video_files"][0]["link"] for v in videos if v["video_files"] and v["video_files"][0]["quality"] == "hd"]
            return random.choice(hd_videos) if hd_videos else None
    except Exception as e:
        print(f"Video search error: {e}")
    return None

def search_pexels_images(query, pexels_api_key):
    headers = {'Authorization': pexels_api_key}
    url = "https://api.pexels.com/v1/search"
    params = {"query": query, "per_page": 5, "orientation": "landscape"}
    try:
        response = requests.get(url, headers=headers, params=params, timeout=10)
        if response.status_code == 200:
            photos = response.json().get("photos", [])
            return random.choice(photos)["src"]["original"] if photos else None
    except Exception as e:
        print(f"Image search error: {e}")
    return None

def search_google_images(query):
    search_url = f"https://www.google.com/search?q={quote(query)}&tbm=isch"
    headers = {"User-Agent": USER_AGENT}
    try:
        response = requests.get(search_url, headers=headers, timeout=10)
        soup = BeautifulSoup(response.text, "html.parser")
        img_urls = [img["src"] for img in soup.find_all("img") if img.get("src", "").startswith("http") and "gstatic" not in img["src"]]
        return random.choice(img_urls[:5]) if img_urls else None
    except Exception as e:
        print(f"Google Images error: {e}")
        return None

def download_image(image_url, filename):
    try:
        response = requests.get(image_url, headers={"User-Agent": USER_AGENT}, stream=True, timeout=15)
        response.raise_for_status()
        with open(filename, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        img = Image.open(filename)
        if img.mode != 'RGB':
            img = img.convert('RGB')
            img.save(filename)
        return filename
    except Exception as e:
        print(f"Image download error: {e}")
        return None

def download_video(video_url, filename):
    try:
        response = requests.get(video_url, stream=True, timeout=30)
        response.raise_for_status()
        with open(filename, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        return filename
    except Exception as e:
        print(f"Video download error: {e}")
        return None

def generate_media(prompt, custom_media=None, video_prob=0.25):
    if custom_media:
        asset_type = "video" if custom_media.endswith(('.mp4', '.avi', '.mov')) else "image"
        return {"path": custom_media, "asset_type": asset_type}
    safe_prompt = re.sub(r'[^\w\s-]', '', prompt).strip().replace(' ', '_')
    if "news" in prompt.lower():
        image_file = os.path.join(TEMP_FOLDER, f"{safe_prompt}_news.jpg")
        image_url = search_google_images(prompt)
        if image_url and download_image(image_url, image_file):
            return {"path": image_file, "asset_type": "image"}
    if random.random() < video_prob:
        video_file = os.path.join(TEMP_FOLDER, f"{safe_prompt}_video.mp4")
        video_url = search_pexels_videos(prompt, PEXELS_API_KEY)
        if video_url and download_video(video_url, video_file):
            return {"path": video_file, "asset_type": "video"}
    image_file = os.path.join(TEMP_FOLDER, f"{safe_prompt}.jpg")
    image_url = search_pexels_images(prompt, PEXELS_API_KEY)
    if image_url and download_image(image_url, image_file):
        return {"path": image_file, "asset_type": "image"}
    return None

def generate_tts(text, voice):
    safe_text = re.sub(r'[^\w\s-]', '', text[:10]).strip().replace(' ', '_')
    file_path = os.path.join(TEMP_FOLDER, f"tts_{safe_text}.wav")
    if os.path.exists(file_path):
        return file_path
    try:
        audio_segments = [audio for _, _, audio in pipeline(text, voice='af_heart', speed=0.9, split_pattern=r'\n+')]
        full_audio = np.concatenate(audio_segments) if len(audio_segments) > 1 else audio_segments[0]
        sf.write(file_path, full_audio, 24000)
        return file_path
    except Exception:
        tts = gTTS(text=text, lang='en')
        mp3_path = os.path.join(TEMP_FOLDER, f"tts_{safe_text}.mp3")
        tts.save(mp3_path)
        AudioSegment.from_mp3(mp3_path).export(file_path, format="wav")
        os.remove(mp3_path)
        return file_path

def apply_kenburns_effect(clip, target_resolution):
    target_w, target_h = target_resolution
    clip_aspect = clip.w / clip.h
    target_aspect = target_w / target_h
    if clip_aspect > target_aspect:
        new_width = int(target_h * clip_aspect)
        clip = clip.resize(width=new_width)
    else:
        new_height = int(target_w / clip_aspect)
        clip = clip.resize(height=new_height)
    clip = clip.resize(zoom=1.15)
    return vfx.crop(clip.fx(vfx.resize, width=target_w * 1.1), width=target_w, height=target_h, x_center=clip.w / 2, y_center=clip.h / 2)

def resize_to_fill(clip, target_resolution):
    target_w, target_h = target_resolution
    clip_aspect = clip.w / clip.h
    target_aspect = target_w / target_h
    if clip_aspect > target_aspect:
        clip = clip.resize(height=target_h)
        crop_amount = (clip.w - target_w) / 2
        clip = clip.crop(x1=crop_amount, x2=clip.w - crop_amount)
    else:
        clip = clip.resize(width=target_w)
        crop_amount = (clip.h - target_h) / 2
        clip = clip.crop(y1=crop_amount, y2=clip.h - crop_amount)
    return clip

def add_background_music(final_video, bgm_path=None, bgm_volume=0.15):
    if bgm_path and os.path.exists(bgm_path):
        bg_music = AudioFileClip(bgm_path)
        if bg_music.duration < final_video.duration:
            loops_needed = math.ceil(final_video.duration / bg_music.duration)
            bg_music = concatenate_audioclips([bg_music] * loops_needed)
        bg_music = bg_music.subclip(0, final_video.duration).volumex(bgm_volume)
        mixed_audio = CompositeAudioClip([final_video.audio, bg_music])
        return final_video.set_audio(mixed_audio)
    return final_video

def create_clip(media_path, asset_type, tts_path, narration_text, target_resolution, subtitles_enabled, font, font_size, outline_width, font_color, outline_color, position, zoom_pan_effect):
    audio_clip = AudioFileClip(tts_path).audio_fadeout(0.2)
    target_duration = audio_clip.duration + 0.2
    if asset_type == "video":
        clip = VideoFileClip(media_path)
        clip = resize_to_fill(clip, target_resolution)
        clip = clip.loop(duration=target_duration) if clip.duration < target_duration else clip.subclip(0, target_duration)
    else:
        clip = ImageClip(media_path).set_duration(target_duration).fadein(0.3).fadeout(0.3)
        if zoom_pan_effect:
            clip = apply_kenburns_effect(clip, target_resolution)
        clip = resize_to_fill(clip, target_resolution)
    if subtitles_enabled and narration_text:
        words = narration_text.split()
        chunks = [' '.join(words[i:i+5]) for i in range(0, len(words), 5)]
        chunk_duration = audio_clip.duration / len(chunks)
        subtitle_clips = []
        y_position = target_resolution[1] * (0.1 if position == "top" else 0.8 if position == "bottom" else 0.5)
        for i, chunk in enumerate(chunks):
            txt_clip = TextClip(
                chunk,
                fontsize=font_size,
                font=font,
                color=font_color,
                stroke_color=outline_color,
                stroke_width=outline_width,
                method='caption',
                align='center',
                size=(target_resolution[0] * 0.8, None)
            ).set_start(i * chunk_duration).set_end((i + 1) * chunk_duration).set_position(('center', y_position))
            subtitle_clips.append(txt_clip)
        clip = CompositeVideoClip([clip] + subtitle_clips)
    return clip.set_audio(audio_clip)

# Main Video Generation Function
def generate_video(resolution, render_speed, video_clip_percent, zoom_pan_effect, bgm_upload, bgm_volume, subtitles_enabled, font, font_size, outline_width, font_color, outline_color, position, *clip_data):
    target_resolution = (1080, 1920) if resolution == "Short (1080x1920)" else (1920, 1080)
    clips = []
    for i in range(0, len(clip_data), 3):
        prompt, narration, custom_media = clip_data[i], clip_data[i+1], clip_data[i+2]
        if prompt.strip() or narration.strip():
            media_asset = generate_media(prompt, custom_media, video_clip_percent / 100.0)
            if not media_asset:
                continue
            tts_path = generate_tts(narration, 'en')
            clip = create_clip(
                media_asset['path'], media_asset['asset_type'], tts_path, narration, target_resolution,
                subtitles_enabled, font, font_size, outline_width, font_color, outline_color, position, zoom_pan_effect
            )
            clips.append(clip)
    if not clips:
        return None
    final_video = concatenate_videoclips(clips, method="compose")
    final_video = add_background_music(final_video, bgm_upload, bgm_volume)
    final_video.write_videofile(OUTPUT_VIDEO_FILENAME, codec='libx264', fps=24, preset=render_speed)
    shutil.rmtree(TEMP_FOLDER)
    return OUTPUT_VIDEO_FILENAME

# Load Clips Function
def load_clips(topic, script):
    raw_script = script.strip() if script.strip() else generate_script(topic)
    clips = parse_script(raw_script)[:MAX_CLIPS]
    updates = [gr.update(value=raw_script)]
    for i in range(MAX_CLIPS):
        if i < len(clips):
            updates.extend([
                gr.update(value=clips[i]["title"]),
                gr.update(value=clips[i]["narration"]),
                gr.update(value=None),
                gr.update(visible=True)
            ])
        else:
            updates.extend([gr.update(value=""), gr.update(value=""), gr.update(value=None), gr.update(visible=False)])
    return updates

# Gradio Interface
with gr.Blocks(title="🚀 Orbit Video Engine") as app:
    with gr.Row():
        # Column 1: Content Input & Script Generation
        with gr.Column():
            gr.Markdown("### 1. Content Input")
            topic_input = gr.Textbox(label="Video Topic", placeholder="e.g., Funny Cat Facts")
            script_input = gr.Textbox(label="Or Paste Full Script", lines=10, placeholder="[Title]\nNarration...")
            generate_script_btn = gr.Button("📝 Generate Script & Load Clips")
            generated_script_display = gr.Textbox(label="Generated Script", interactive=False)

        # Column 2: Clip Editor
        with gr.Column():
            gr.Markdown("### 2. Edit Clips")
            gr.Markdown("Modify prompts, narration, and upload custom media for each clip.")
            clip_rows, prompts, narrations, custom_medias = [], [], [], []
            for i in range(MAX_CLIPS):
                with gr.Row(visible=False) as row:
                    prompt = gr.Textbox(label="Visual Prompt")
                    narration = gr.Textbox(label="Narration", lines=3)
                    custom_media = gr.File(label="Upload Custom Media (Image/Video)", file_types=["image", "video"])
                clip_rows.append(row)
                prompts.append(prompt)
                narrations.append(narration)
                custom_medias.append(custom_media)

        # Column 3: Settings & Output
        with gr.Column():
            gr.Markdown("### 3. Video Settings")
            resolution = gr.Radio(["Short (1080x1920)", "Full HD (1920x1080)"], label="Resolution", value="Full HD (1920x1080)")
            render_speed = gr.Dropdown(["ultrafast", "veryfast", "fast", "medium", "slow", "veryslow"], label="Render Speed", value="fast")
            video_clip_percent = gr.Slider(0, 100, value=25, label="Video Clip Percentage")
            zoom_pan_effect = gr.Checkbox(label="Add Zoom/Pan Effect (Images)", value=True)
            with gr.Accordion("Background Music", open=False):
                bgm_upload = gr.Audio(label="Upload Background Music", type="filepath")
                bgm_volume = gr.Slider(0.0, 1.0, value=0.15, label="BGM Volume")
            with gr.Accordion("Subtitle Settings", open=True):
                subtitles_enabled = gr.Checkbox(label="Enable Subtitles", value=True)
                font = gr.Dropdown(["Impact", "Arial", "Times New Roman"], label="Font", value="Arial")
                font_size = gr.Number(label="Font Size", value=45)
                outline_width = gr.Number(label="Outline Width", value=2)
                font_color = gr.ColorPicker(label="Font Color", value="#FFFFFF")
                outline_color = gr.ColorPicker(label="Outline Color", value="#000000")
                position = gr.Radio(["top", "center", "bottom"], label="Position", value="bottom")
            generate_video_btn = gr.Button("🎬 Generate Video")
            gr.Markdown("### 4. Output")
            video_output = gr.Video(label="Generated Video")

    # Event Handlers
    generate_script_btn.click(
        load_clips,
        inputs=[topic_input, script_input],
        outputs=[generated_script_display] + prompts + narrations + custom_medias + clip_rows
    )
    generate_video_btn.click(
        generate_video,
        inputs=[resolution, render_speed, video_clip_percent, zoom_pan_effect, bgm_upload, bgm_volume, subtitles_enabled, font, font_size, outline_width, font_color, outline_color, position] + prompts + narrations + custom_medias,
        outputs=video_output
    )

app.launch(share=True)