import gradio as gr
import librosa
import numpy as np
import tempfile
import os
import random
from moviepy.editor import (
    VideoFileClip,
    ImageClip,
    AudioFileClip,
    concatenate_videoclips,
    vfx
)

# -------------------------
# Beat + energy detection
# -------------------------
def analyze_audio(audio_path):
    y, sr = librosa.load(audio_path)
    tempo, beats = librosa.beat.beat_track(y=y, sr=sr)
    beat_times = librosa.frames_to_time(beats, sr=sr)

    onset_env = librosa.onset.onset_strength(y=y, sr=sr)
    energy = onset_env / np.max(onset_env)

    return beat_times, energy

# -------------------------
# Template rules
# -------------------------
TEMPLATES = {
    "Auto": {
        "min_clip": 0.3,
        "max_clip": 0.8,
        "zoom_prob": 0.4,
        "shake_prob": 0.2
    },
    "Hype": {
        "min_clip": 0.2,
        "max_clip": 0.5,
        "zoom_prob": 0.6,
        "shake_prob": 0.4
    },
    "Sad": {
        "min_clip": 0.6,
        "max_clip": 1.2,
        "zoom_prob": 0.2,
        "shake_prob": 0.0
    }
}

# -------------------------
# Edit generator
# -------------------------
def generate_edit(song, media_files, template):
    if song is None or not media_files:
        return None

    rules = TEMPLATES[template]

    with tempfile.TemporaryDirectory() as tmp:
        song_path = os.path.join(tmp, "song.mp3")
        with open(song_path, "wb") as f:
            f.write(song)

        beats, energy = analyze_audio(song_path)
        audio = AudioFileClip(song_path)

        clips = []
        media_index = 0
        t = 0

        while t < audio.duration:
            duration = random.uniform(
                rules["min_clip"], rules["max_clip"]
            )

            media = media_files[media_index % len(media_files)]
            media_index += 1

            media_path = os.path.join(tmp, os.path.basename(media.name))
            with open(media_path, "wb") as f:
                f.write(media.read())

            if media_path.lower().endswith((".mp4", ".mov", ".webm")):
                base = VideoFileClip(media_path)
                clip = base.subclip(0, min(duration, base.duration))
            else:
                clip = ImageClip(media_path).set_duration(duration)

            # Effects
            if random.random() < rules["zoom_prob"]:
                clip = clip.fx(vfx.resize, lambda t: 1 + 0.05 * t)

            if random.random() < rules["shake_prob"]:
                clip = clip.fx(vfx.crop, x1=5, y1=5, x2=5, y2=5)

            clip = clip.resize(height=1080).set_position("center")
            clips.append(clip)

            t += duration

        final = concatenate_videoclips(
            clips, method="compose", padding=-0.05
        )

        final = final.set_audio(audio)
        final = final.resize((1080, 1920))

        output = os.path.join(tmp, "edit.mp4")
        final.write_videofile(
            output,
            fps=30,
            codec="libx264",
            audio_codec="aac",
            threads=2,
            logger=None
        )

        return output

# -------------------------
# UI
# -------------------------
with gr.Blocks(title="EditForge AI Advanced") as demo:
    gr.Markdown("## 🎬 EditForge AI (Advanced)")

    song = gr.File(label="Upload Song", type="binary")
    media = gr.Files(label="Upload Images / Videos")

    template = gr.Dropdown(
        choices=["Auto", "Hype", "Sad"],
        value="Auto",
        label="Edit Style"
    )

    output = gr.Video(label="Final Edit")

    btn = gr.Button("🚀 Generate Edit")

    btn.click(
        fn=generate_edit,
        inputs=[song, media, template],
        outputs=output
    )

demo.launch(server_name="0.0.0.0", server_port=7860)