import streamlit as st
import numpy as np
from PIL import Image
import torch
import tempfile
import os
import soundfile as sf
import imageio.v2 as imageio

from transformers import AutoProcessor, BlipForConditionalGeneration, MusicgenForConditionalGeneration

import moviepy.editor as mpy


st.title("Video → Scene Sound Generator (Optimized)")

num_frames = st.slider("Frames to sample", 1, 5, 2)
mix_audio = st.checkbox("Mix original audio", False)


# -----------------------------
# SAFE PROMPT COMPRESSION
# -----------------------------
def build_prompt(captions):
    captions = [c.lower().strip() for c in captions if c]

    if not captions:
        return "ambient background sound"

    core = captions[0]

    # aggressively compress noise
    keywords = []
    for c in captions:
        for word in ["street", "car", "person", "room", "forest", "dog", "run", "walk", "talk"]:
            if word in c:
                keywords.append(word)

    keywords = list(set(keywords))

    return f"{core}. scene contains {' '.join(keywords)}. ambient cinematic sound."


# -----------------------------
# MODEL LOADING (CACHED)
# -----------------------------
@st.cache_resource
def load_blip():
    processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
    model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

    if torch.cuda.is_available():
        model = model.to("cuda").half()

    return processor, model


@st.cache_resource
def load_musicgen():
    processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
    model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")

    if torch.cuda.is_available():
        model = model.to("cuda").half()

    return processor, model


blip_processor, blip_model = load_blip()
mg_processor, mg_model = load_musicgen()


# -----------------------------
# SAFE FRAME SAMPLING (FIXED)
# -----------------------------
def sample_frames(video_path, n):
    reader = imageio.get_reader(video_path, "ffmpeg")

    meta = reader.get_meta_data()
    duration = meta.get("duration", 10)

    timestamps = np.linspace(0, duration, n + 2)[1:-1]

    frames = []

    for t in timestamps:
        try:
            frame = reader.get_data(int(t * meta.get("fps", 24)))
            frames.append(Image.fromarray(frame))
        except:
            continue

    reader.close()
    return frames


# -----------------------------
# AUDIO NORMALIZATION SAFE
# -----------------------------
def safe_normalize(audio):
    max_val = np.max(np.abs(audio))
    if max_val < 1e-8:
        return audio
    return (audio / max_val) * 0.8


# -----------------------------
# UPLOAD
# -----------------------------
file = st.file_uploader("Upload MP4", type=["mp4"])

if file:

    temp_video = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    temp_video.write(file.read())
    temp_video.close()

    try:
        frames = sample_frames(temp_video.name, num_frames)

        captions = []

        for f in frames:
            inputs = blip_processor(images=f, return_tensors="pt")

            if torch.cuda.is_available():
                inputs = {k: v.to("cuda") for k, v in inputs.items()}

            out = blip_model.generate(**inputs)
            cap = blip_processor.decode(out[0], skip_special_tokens=True)

            captions.append(cap)

        prompt = build_prompt(captions)
        st.write("Prompt:", prompt)

        # -----------------------------
        # AUDIO GENERATION
        # -----------------------------
        inputs = mg_processor(text=[prompt], return_tensors="pt")

        if torch.cuda.is_available():
            inputs = {k: v.to("cuda") for k, v in inputs.items()}

        audio = mg_model.generate(
            **inputs,
            max_new_tokens=256,
            do_sample=True
        )

        audio = audio[0].cpu().numpy()
        audio = safe_normalize(audio)

        wav_path = tempfile.NamedTemporaryFile(delete=False, suffix=".wav").name
        sf.write(wav_path, audio, 32000)

        # -----------------------------
        # VIDEO + AUDIO MERGE (SIMPLIFIED)
        # -----------------------------
        video = mpy.VideoFileClip(temp_video.name)
        audio_clip = mpy.AudioFileClip(wav_path)

        if audio_clip.duration > video.duration:
            audio_clip = audio_clip.subclip(0, video.duration)

        if mix_audio and video.audio:
            final_audio = video.audio.volumex(0.4).fx(lambda a: a) + audio_clip.volumex(0.6)
        else:
            final_audio = audio_clip

        final = video.set_audio(final_audio)

        out_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name

        final.write_videofile(
            out_path,
            codec="libx264",
            audio_codec="aac",
            fps=24,
            bitrate="4000k",
            preset="ultrafast"
        )

        st.video(out_path)

        with open(out_path, "rb") as f:
            st.download_button("Download", f, file_name="output.mp4")

    finally:
        for p in [temp_video.name, wav_path, out_path]:
            if p and os.path.exists(p):
                os.remove(p)