app.py

import os
import numpy as np
import torch
import typing
import torchaudio
import streamlit as st
import subprocess
import json
import requests
import gc
import math
import uuid
import re
import time
from urllib.parse import quote
from datetime import timedelta

# --- CRITICAL ENVIRONMENT FIXES ---
if os.environ.get("OMP_NUM_THREADS", "").endswith("m"):
    os.environ["OMP_NUM_THREADS"] = "1"

try:
    if "ffmpeg" in torchaudio.list_audio_backends():
        torchaudio.set_audio_backend("ffmpeg")
except Exception:
    pass

if not hasattr(torch.load, "_is_patched"):
    _original_torch_load = torch.load

    def patched_torch_load(*args, **kwargs):
        kwargs['weights_only'] = False 
        return _original_torch_load(*args, **kwargs)
    
    patched_torch_load._is_patched = True
    torch.load = patched_torch_load

try:
    safe_list = [typing.Any, torch.nn.modules.container.ModuleList, np.dtype]
    if hasattr(np, '_core') and hasattr(np._core, 'multiarray'):
        safe_list.append(np._core.multiarray.scalar)
    elif hasattr(np, 'core') and hasattr(np.core, 'multiarray'):
        safe_list.append(np.core.multiarray.scalar)
    
    try:
        from omegaconf.listconfig import ListConfig
        from omegaconf.dictconfig import DictConfig
        from omegaconf.base import ContainerMetadata, Metadata, Node
        safe_list.extend([ListConfig, DictConfig, ContainerMetadata, Metadata, Node])
    except ImportError:
        pass

    try:
        from pyannote.audio.core.task import Specifications, Problem, Resolution
        from pyannote.audio.core.model import Model
        from pyannote.audio.pipelines.speaker_diarization import SpeakerDiarization
        safe_list.extend([Specifications, Problem, Resolution, Model, SpeakerDiarization])
    except ImportError:
        pass

    torch.serialization.add_safe_globals(safe_list)
except Exception as e:
    print(f"Safe Globals Warning: {e}")

if not hasattr(np, 'NaN'):
    np.NaN = np.nan

import whisperx

# --- Configuration ---
HARDCODED_HF_TOKEN = "PASTE_YOUR_HF_TOKEN_HERE" 
HARDCODED_GEMINI_KEY = ""

ENV_HF_TOKEN = os.environ.get("HF_TOKEN", "")
ENV_GEMINI_KEY = os.environ.get("GEMINI_API_KEY", "")

ACTIVE_HF_TOKEN = ENV_HF_TOKEN if ENV_HF_TOKEN else HARDCODED_HF_TOKEN
ACTIVE_GEMINI_KEY = ENV_GEMINI_KEY if ENV_GEMINI_KEY else HARDCODED_GEMINI_KEY

# --- HELPER FUNCTIONS ---

def clean_json_response(text):
    """Strips markdown code fences from Gemini output to prevent JSON errors."""
    try:
        pattern = r"`{3}(?:json)?\s*(.*?)`{3}"
        match = re.search(pattern, text, re.DOTALL)
        if match:
            return match.group(1).strip()
        return text.strip()
    except Exception:
        return text

def escape_xml(text):
    """Safely escapes special characters for XML to prevent DOM parser errors."""
    if not text:
        return ""
    text = str(text)
    text = text.replace("&", "&")
    text = text.replace("<", "&lt;")
    text = text.replace(">", "&gt;")
    text = text.replace('"', "&quot;")
    text = text.replace("'", "&apos;")
    return text

def timecode_to_frames(tc, fps):
    """Converts HH:MM:SS:FF to absolute integer frames based on timebase."""
    if not tc or not re.match(r"\d{2}:\d{2}:\d{2}[:\.]\d{2}", tc):
        return 0
    parts = re.split(r'[:\.]', tc)
    h, m, s, f = map(int, parts)
    timebase = int(round(fps)) # 29.97 becomes 30, 23.98 becomes 24
    return (h * 3600 * timebase) + (m * 60 * timebase) + (s * timebase) + f

def frames_to_timecode(frames, fps):
    """Converts absolute integer frames back to HH:MM:SS:FF (NDF)."""
    timebase = int(round(fps))
    hours = frames // (3600 * timebase)
    minutes = (frames // (60 * timebase)) % 60
    secs = (frames // timebase) % 60
    f = frames % timebase
    return f"{hours:02}:{minutes:02}:{secs:02}:{f:02}"

def seconds_to_frames(seconds, fps):
    """Converts real seconds to integer frames."""
    return int(round(seconds * fps))

def generate_cmx_edl(edl_title, segments, source_name, fps=25, start_offset_frames=0):
    """Constructs a CMX 3600 formatted EDL with source offset."""
    edl_lines = [f"TITLE: {edl_title}", "FCM: NON-DROP FRAME\n"]
    rec_start_frames = 0
    reel_id = source_name.replace(" ", "_")[:8] # EDL reel IDs are traditionally short max 8 chars
    
    for i, seg in enumerate(segments, 1):
        seg_type = seg.get('type', 'clip')
        is_vo = seg_type == 'vo' or ('text' in seg and seg_type != 'graphic')
        is_graphic = seg_type == 'graphic'
        
        if is_vo or is_graphic:
            text_content = seg.get('text', 'Placeholder')
            default_dur = max(1.0, len(text_content.split()) / 2.5) if is_vo else 4.0
            duration_sec = float(seg.get('duration', default_dur))
            duration_frames = seconds_to_frames(duration_sec, fps)
            src_in_frames = 0
            src_out_frames = duration_frames
            current_source = "GEN_VO" if is_vo else "GEN_GFX"
            note_text = f"{'VO SCRIPT' if is_vo else 'GRAPHIC'}: {text_content}"
        else:
            src_start_val = float(seg.get('src_start', seg.get('start', 0.0)))
            src_end_val = float(seg.get('src_end', seg.get('end', 0.0)))
            src_in_frames = seconds_to_frames(src_start_val, fps) + start_offset_frames
            src_out_frames = seconds_to_frames(src_end_val, fps) + start_offset_frames
            duration_frames = src_out_frames - src_in_frames
            current_source = reel_id
            note_text = seg.get('note', 'Clip')
            
        src_in = frames_to_timecode(src_in_frames, fps)
        src_out = frames_to_timecode(src_out_frames, fps)
        
        rec_in = frames_to_timecode(rec_start_frames, fps)
        rec_out = frames_to_timecode(rec_start_frames + duration_frames, fps)
        
        edl_lines.append(f"{i:03}  {current_source:8} V     C        {src_in} {src_out} {rec_in} {rec_out}")
        if not is_vo and not is_graphic:
            edl_lines.append(f"* FROM CLIP NAME: {source_name}")
        edl_lines.append(f"* {note_text}\n")
        
        gap_frames = seconds_to_frames(float(seg.get('gap', 0.0)), fps)
        rec_start_frames += duration_frames + gap_frames
        
    return "\n".join(edl_lines)

def generate_xml(sequence_name, segments, source_name, fps=25, start_offset_frames=0):
    """Constructs a Robust XML with source timecode offset for Premiere AND Resolve."""
    master_id = "masterfile-1"
    timebase = int(round(fps))
    is_ntsc = "TRUE" if fps % 1 != 0 else "FALSE"
    
    start_tc_string = frames_to_timecode(start_offset_frames, fps)

    encoded_name = quote(source_name)
    path_url = f"file://localhost/{encoded_name}" 

    lines = []
    lines.append('<?xml version="1.0" encoding="UTF-8"?>')
    lines.append('<!DOCTYPE xmeml>')
    lines.append('<xmeml version="4">')
    lines.append('<sequence>')
    lines.append(f'\t<name>{escape_xml(sequence_name)}</name>')
    lines.append('\t<rate>')
    lines.append(f'\t\t<timebase>{timebase}</timebase>')
    lines.append(f'\t\t<ntsc>{is_ntsc}</ntsc>') 
    lines.append('\t</rate>')
    lines.append('\t<media>')
    
    # --- VIDEO TRACK ---
    lines.append('\t\t<video>')
    lines.append('\t\t\t<format>')
    lines.append('\t\t\t\t<samplecharacteristics>')
    lines.append(f'\t\t\t\t\t<rate><timebase>{timebase}</timebase></rate>')
    lines.append('\t\t\t\t\t<width>1920</width>')
    lines.append('\t\t\t\t\t<height>1080</height>')
    lines.append('\t\t\t\t\t<anamorphic>FALSE</anamorphic>')
    lines.append('\t\t\t\t\t<pixelaspectratio>square</pixelaspectratio>')
    lines.append('\t\t\t\t</samplecharacteristics>')
    lines.append('\t\t\t</format>')
    lines.append('\t\t\t<track>')

    file_defined = False
    vo_file_defined = False
    gfx_file_defined = False
    timeline_head_frames = 0
    
    for i, seg in enumerate(segments, 1):
        seg_type = seg.get('type', 'clip')
        is_vo = seg_type == 'vo' or ('text' in seg and seg_type != 'graphic')
        is_graphic = seg_type == 'graphic'
        
        if is_vo or is_graphic:
            text_content = seg.get('text', 'Placeholder')
            default_dur = max(1.0, len(text_content.split()) / 2.5) if is_vo else 4.0
            duration_sec = float(seg.get('duration', default_dur))
            duration_frames = seconds_to_frames(duration_sec, fps)
            src_in_frames = 0
            src_out_frames = duration_frames
            master_id_to_use = "masterfile-vo" if is_vo else "masterfile-gfx"
            source_name_to_use = "Generated_VO_Placeholder" if is_vo else "Generated_Graphic_Placeholder"
            prefix = "VO SCRIPT" if is_vo else "GRAPHIC CARD"
            clip_note = f"{prefix}: {text_content}"
            clip_name_disp = (clip_note[:75] + '..') if len(clip_note) > 75 else clip_note
        else:
            src_start_val = float(seg.get('src_start', seg.get('start', 0.0)))
            src_end_val = float(seg.get('src_end', seg.get('end', 0.0)))
            
            # FIX: Removed the '+ start_offset_frames' from the video track so it matches the audio track!
            src_in_frames = seconds_to_frames(src_start_val, fps)
            src_out_frames = seconds_to_frames(src_end_val, fps)
            
            duration_frames = src_out_frames - src_in_frames
            master_id_to_use = master_id
            source_name_to_use = source_name
            raw_note = seg.get('note', 'Junior Editor Selection')
            clip_note = raw_note
            clip_name_disp = source_name
        
        tl_start = timeline_head_frames
        tl_end = tl_start + duration_frames

        lines.append(f'\t\t\t\t<clipitem id="clipitem-v-{i}">')
        lines.append(f'\t\t\t\t\t<name>{escape_xml(clip_name_disp)}</name>')
        lines.append(f'\t\t\t\t\t<duration>8640000</duration>')
        lines.append(f'\t\t\t\t\t<rate><timebase>{timebase}</timebase></rate>')
        lines.append(f'\t\t\t\t\t<start>{tl_start}</start>')
        lines.append(f'\t\t\t\t\t<end>{tl_end}</end>')
        lines.append(f'\t\t\t\t\t<in>{src_in_frames}</in>')
        lines.append(f'\t\t\t\t\t<out>{src_out_frames}</out>')
        lines.append(f'\t\t\t\t\t<masterclipid>{master_id_to_use}</masterclipid>')
        
        # --- ROBUST FILE DEFINITIONS ---
        if is_vo:
            if not vo_file_defined:
                lines.append(f'\t\t\t\t\t<file id="{master_id_to_use}">')
                lines.append(f'\t\t\t\t\t\t<name>{escape_xml(source_name_to_use)}</name>')
                lines.append(f'\t\t\t\t\t\t<pathurl>file://localhost/{escape_xml(source_name_to_use)}.wav</pathurl>') 
                lines.append(f'\t\t\t\t\t\t<rate><timebase>{timebase}</timebase></rate>')
                lines.append(f'\t\t\t\t\t\t<duration>8640000</duration>') 
                lines.append(f'\t\t\t\t\t\t<timecode>')
                lines.append(f'\t\t\t\t\t\t\t<rate><timebase>{timebase}</timebase></rate>')
                lines.append(f'\t\t\t\t\t\t\t<string>00:00:00:00</string>') 
                lines.append(f'\t\t\t\t\t\t\t<frame>0</frame>')
                lines.append(f'\t\t\t\t\t\t\t<displayformat>NDF</displayformat>')
                lines.append(f'\t\t\t\t\t\t</timecode>')
                lines.append('\t\t\t\t\t\t<media>')
                lines.append('\t\t\t\t\t\t\t<video><samplecharacteristics><width>1920</width><height>1080</height></samplecharacteristics></video>')
                lines.append('\t\t\t\t\t\t\t<audio><samplecharacteristics><depth>16</depth><samplerate>48000</samplerate></samplecharacteristics><channelcount>2</channelcount></audio>')
                lines.append('\t\t\t\t\t\t</media>')
                lines.append('\t\t\t\t\t</file>')
                vo_file_defined = True
            else:
                lines.append(f'\t\t\t\t\t<file id="{master_id_to_use}"/>')
        elif is_graphic:
            if not gfx_file_defined:
                lines.append(f'\t\t\t\t\t<file id="{master_id_to_use}">')
                lines.append(f'\t\t\t\t\t\t<name>{escape_xml(source_name_to_use)}</name>')
                lines.append(f'\t\t\t\t\t\t<pathurl>file://localhost/{escape_xml(source_name_to_use)}.mov</pathurl>') 
                lines.append(f'\t\t\t\t\t\t<rate><timebase>{timebase}</timebase></rate>')
                lines.append(f'\t\t\t\t\t\t<duration>8640000</duration>') 
                lines.append(f'\t\t\t\t\t\t<timecode>')
                lines.append(f'\t\t\t\t\t\t\t<rate><timebase>{timebase}</timebase></rate>')
                lines.append(f'\t\t\t\t\t\t\t<string>00:00:00:00</string>') 
                lines.append(f'\t\t\t\t\t\t\t<frame>0</frame>')
                lines.append(f'\t\t\t\t\t\t\t<displayformat>NDF</displayformat>')
                lines.append(f'\t\t\t\t\t\t</timecode>')
                lines.append('\t\t\t\t\t\t<media>')
                lines.append('\t\t\t\t\t\t\t<video><samplecharacteristics><width>1920</width><height>1080</height></samplecharacteristics></video>')
                # Intentionally omitting audio block for graphics placeholder to keep NLE track structure clean
                lines.append('\t\t\t\t\t\t</media>')
                lines.append('\t\t\t\t\t</file>')
                gfx_file_defined = True
            else:
                lines.append(f'\t\t\t\t\t<file id="{master_id_to_use}"/>')
        else:
            if not file_defined:
                lines.append(f'\t\t\t\t\t<file id="{master_id_to_use}">')
                lines.append(f'\t\t\t\t\t\t<name>{escape_xml(source_name_to_use)}</name>')
                lines.append(f'\t\t\t\t\t\t<pathurl>{path_url}</pathurl>') 
                lines.append(f'\t\t\t\t\t\t<rate><timebase>{timebase}</timebase></rate>')
                lines.append(f'\t\t\t\t\t\t<duration>8640000</duration>') 
                lines.append(f'\t\t\t\t\t\t<timecode>')
                lines.append(f'\t\t\t\t\t\t\t<rate><timebase>{timebase}</timebase></rate>')
                lines.append(f'\t\t\t\t\t\t\t<string>{start_tc_string}</string>') 
                lines.append(f'\t\t\t\t\t\t\t<frame>{start_offset_frames}</frame>')
                lines.append(f'\t\t\t\t\t\t\t<displayformat>NDF</displayformat>')
                lines.append(f'\t\t\t\t\t\t</timecode>')
                lines.append('\t\t\t\t\t\t<media>')
                lines.append('\t\t\t\t\t\t\t<video><samplecharacteristics><width>1920</width><height>1080</height></samplecharacteristics></video>')
                lines.append('\t\t\t\t\t\t\t<audio><samplecharacteristics><depth>16</depth><samplerate>48000</samplerate></samplecharacteristics><channelcount>2</channelcount></audio>')
                lines.append('\t\t\t\t\t\t</media>')
                lines.append('\t\t\t\t\t</file>')
                file_defined = True
            else:
                lines.append(f'\t\t\t\t\t<file id="{master_id_to_use}"/>')
            
        lines.append('\t\t\t\t\t<marker>')
        lines.append(f'\t\t\t\t\t\t<name>{escape_xml(clip_note)}</name>')
        lines.append(f'\t\t\t\t\t\t<in>{src_in_frames}</in>')
        
        # Extends the marker duration to span the entire clip for VO and Graphics
        if is_vo or is_graphic:
            lines.append(f'\t\t\t\t\t\t<out>{src_out_frames}</out>')
        else:
            lines.append(f'\t\t\t\t\t\t<out>{src_in_frames}</out>')
            
        lines.append('\t\t\t\t\t</marker>')

        lines.append('\t\t\t\t</clipitem>')
        
        gap_frames = seconds_to_frames(float(seg.get('gap', 0.0)), fps)
        timeline_head_frames = tl_end + gap_frames

    lines.append('\t\t\t</track>')
    lines.append('\t\t</video>')
    
    # --- AUDIO TRACK ---
    lines.append('\t\t<audio>')
    lines.append('\t\t\t<track>')

    timeline_head_frames = 0
    for i, seg in enumerate(segments, 1):
        seg_type = seg.get('type', 'clip')
        is_vo = seg_type == 'vo' or ('text' in seg and seg_type != 'graphic')
        is_graphic = seg_type == 'graphic'
        
        if is_graphic:
            # Graphics have no audio. We just advance the playhead so the timeline stays perfectly in sync!
            duration_sec = float(seg.get('duration', 4.0))
            duration_frames = seconds_to_frames(duration_sec, fps)
            gap_frames = seconds_to_frames(float(seg.get('gap', 0.0)), fps)
            timeline_head_frames += duration_frames + gap_frames
            continue
            
        if is_vo:
            vo_text = seg.get('text', 'VO')
            duration_sec = float(seg.get('duration', max(1.0, len(vo_text.split()) / 2.5)))
            duration_frames = seconds_to_frames(duration_sec, fps)
            src_in_frames = 0
            src_out_frames = duration_frames
            master_id_to_use = "masterfile-vo"
            source_name_to_use = "Generated_VO_Placeholder"
        else:
            src_start_val = float(seg.get('src_start', seg.get('start', 0.0)))
            src_end_val = float(seg.get('src_end', seg.get('end', 0.0)))
            src_in_frames = seconds_to_frames(src_start_val, fps)
            src_out_frames = seconds_to_frames(src_end_val, fps)
            duration_frames = src_out_frames - src_in_frames
            master_id_to_use = master_id
            source_name_to_use = source_name
            
        tl_start = timeline_head_frames
        tl_end = tl_start + duration_frames

        lines.append(f'\t\t\t\t<clipitem id="clipitem-a-{i}">')
        lines.append(f'\t\t\t\t\t<name>{escape_xml(source_name_to_use)}</name>')
        lines.append(f'\t\t\t\t\t<masterclipid>{master_id_to_use}</masterclipid>')
        lines.append(f'\t\t\t\t\t<duration>8640000</duration>')
        lines.append(f'\t\t\t\t\t<rate><timebase>{timebase}</timebase></rate>')
        lines.append(f'\t\t\t\t\t<start>{tl_start}</start>')
        lines.append(f'\t\t\t\t\t<end>{tl_end}</end>')
        lines.append(f'\t\t\t\t\t<in>{src_in_frames}</in>')
        lines.append(f'\t\t\t\t\t<out>{src_out_frames}</out>')
        lines.append(f'\t\t\t\t\t<file id="{master_id_to_use}"/>')
        lines.append('\t\t\t\t\t<sourcetrack><mediatype>audio</mediatype><trackindex>1</trackindex></sourcetrack>')
        lines.append('\t\t\t\t</clipitem>')
        timeline_head_frames = tl_end + seconds_to_frames(float(seg.get('gap', 0.0)), fps)

    lines.append('\t\t\t</track>')
    lines.append('\t\t</audio>')
    lines.append('\t</media>')
    lines.append('</sequence>')
    lines.append('</xmeml>')
    
    return "\n".join(lines)

def generate_transcript_txt(sequence_name, segments, fps=25, start_offset_frames=0, transcript_data=None):
    """Generates a human-readable text document of the edit decisions and exact transcription."""
    lines = [f"Sequence Transcript: {sequence_name}", "=" * 50, ""]
    
    for i, seg in enumerate(segments, 1):
        seg_type = seg.get('type', 'clip')
        is_vo = seg_type == 'vo' or ('text' in seg and seg_type != 'graphic')
        is_graphic = seg_type == 'graphic'
        
        if is_vo:
            vo_text = seg.get('text', 'VO')
            duration = float(seg.get('duration', max(1.0, len(vo_text.split()) / 2.5)))
            lines.append(f"Clip {i:02}  |  [GENERATED VOICE-OVER]  |  Est. Duration: {duration:.1f}s")
            lines.append(f"Script: \"{vo_text}\"")
        elif is_graphic:
            gfx_text = seg.get('text', 'Graphic Text')
            duration = float(seg.get('duration', 4.0))
            lines.append(f"Clip {i:02}  |  [GRAPHIC CARD]  |  Est. Duration: {duration:.1f}s")
            lines.append(f"Text: \"{gfx_text}\"")
        else:
            src_start_val = float(seg.get('src_start', seg.get('start', 0.0)))
            src_end_val = float(seg.get('src_end', seg.get('end', 0.0)))
            src_in_frames = seconds_to_frames(src_start_val, fps) + start_offset_frames
            src_out_frames = seconds_to_frames(src_end_val, fps) + start_offset_frames
            
            src_in = frames_to_timecode(src_in_frames, fps)
            src_out = frames_to_timecode(src_out_frames, fps)
            note = seg.get('note', 'No note provided.')
            
            actual_text = ""
            if transcript_data:
                clip_words = []
                for t_seg in transcript_data:
                    if max(src_start_val, t_seg['start']) < min(src_end_val, t_seg['end']):
                        clip_words.append(t_seg['text'])
                actual_text = " ".join(clip_words).strip()

            lines.append(f"Clip {i:02}  |  IN: {src_in}  -->  OUT: {src_out}")
            lines.append(f"AI Summary: {note}")
            if actual_text:
                lines.append(f"Exact Audio Text: \"{actual_text}\"")
            
        lines.append("") 
        
    return "\n".join(lines)


def call_gemini_for_edl(transcript_data, story_prompt, api_key, previous_edit=None):
    if not api_key:
        st.error("Gemini API Key is missing.")
        return None
        
    system_prompt = (
        "You are an expert Documentary Senior Editor. Use the provided transcript JSON "
        "(which includes Speaker IDs and word-level timestamps) to create a condensed story.\n\n"
        "NEW CAPABILITIES:\n"
        "1. You can generate Voice-Over (VO) to bridge gaps, introduce topics, or summarize.\n"
        "2. You can generate Graphic Cards (Title Cards) to display text on screen (e.g., location, date, or chapter title).\n"
        "CRITICAL: ONLY generate 'vo' or 'graphic' segments if the user explicitly requests voice over, narration, or graphic/title cards in their creative brief. Otherwise, ONLY use 'clip' segments.\n\n"
        "Output ONLY a valid JSON array of segments. Every segment MUST be a 'clip', a 'vo', or a 'graphic'.\n\n"
        "For 'clip' segments (extracting from the subject):\n"
        "{\"type\": \"clip\", \"src_start\": 12.5, \"src_end\": 25.0, \"note\": \"Subject talks about X\", \"gap\": 0.0}\n"
        "- IGNORE ALL INTERVIEWER COMMENTS.\n"
        "- REMOVE FLUFF: Delete 'um', 'ah', repeats, and irrelevant filler.\n"
        "- TIMESTAMP INTEGRITY: Use only the exact word-level start and end times from the data.\n\n"
        "For 'vo' segments (generating new voice-over):\n"
        "{\"type\": \"vo\", \"text\": \"The journey began years ago...\", \"duration\": 3.5, \"gap\": 0.0}\n"
        "- DURATION: Estimate duration accurately based on reading speed (approx 2.5 words per second).\n"
        "- Keep VO concise and in the style of a documentary narrator.\n\n"
        "For 'graphic' segments (generating text on screen):\n"
        "{\"type\": \"graphic\", \"text\": \"Chapter 1: The Beginning\", \"duration\": 4.0, \"gap\": 0.0}\n\n"
        "PACING: Group related clips. Add a 'gap' (in seconds) between distinct ideas."
    )
    
    if previous_edit:
        prompt_text = (
            f"Transcript Data:\n{json.dumps(transcript_data)}\n\n"
            f"Previous Edit You Generated:\n{json.dumps(previous_edit)}\n\n"
            f"Director's Note for Revision: {story_prompt}\n\n"
            "Please output a NEW JSON array of segments applying these requested changes to the previous edit."
        )
    else:
        prompt_text = f"Creative Brief: {story_prompt}\n\nTranscript Data:\n{json.dumps(transcript_data)}"
    
    payload = {
        "contents": [
            {
                "role": "user",
                "parts": [{"text": prompt_text}]
            }
        ],
        "systemInstruction": {
            "parts": [{"text": system_prompt}]
        },
        "generationConfig": {
            "responseMimeType": "application/json"
        }
    }
    
    headers = {
        "Content-Type": "application/json"
    }
    
    # Prioritizes Pro models, falls back to Flash. All "lite" versions explicitly removed.
    models_to_try = [
        "gemini-3.1-pro-preview",
        "gemini-3-flash-preview",
        "gemini-2.5-pro",
        "gemini-2.5-flash",
        "gemini-2.0-flash"
    ]
    
    endpoints_to_try = ["v1beta", "v1"]
    last_error = ""
    
    for model_name in models_to_try:
        for api_version in endpoints_to_try:
            url = f"https://generativelanguage.googleapis.com/{api_version}/models/{model_name}:generateContent?key={api_key}"
            
            try:
                res = requests.post(url, json=payload, headers=headers)
                
                # Success
                if res.status_code == 200:
                    result_json = res.json()
                    raw_text = result_json['candidates'][0]['content']['parts'][0]['text']
                    cleaned_text = clean_json_response(raw_text)
                    return json.loads(cleaned_text), model_name, api_version
                    
                # Handle High Demand (503) or Rate Limits (429)
                elif res.status_code in [503, 429]:
                    last_error = f"Model {model_name} busy ({res.status_code}) on {api_version}."
                    st.warning(f"Debug: {last_error} Skipping to next model...")
                    break 
                    
                # Handle endpoints that don't exist (404) or API rejections (400)
                else:
                    try:
                        error_details = res.json().get('error', {}).get('message', res.text)
                    except Exception:
                        error_details = res.text
                    last_error = f"Model {model_name} rejected request ({res.status_code}) on {api_version}: {error_details}"
                    st.warning(f"Debug: {last_error} Trying fallback endpoint/model...")
                    continue 
                    
            except json.JSONDecodeError:
                last_error = f"Model {model_name} returned invalid JSON on {api_version}."
                st.warning(f"Debug: {last_error} Trying fallback endpoint/model...")
                continue
            except requests.exceptions.RequestException as e:
                last_error = f"Network error connecting to {model_name} on {api_version}: {e}"
                st.warning(f"Debug: {last_error} Trying fallback endpoint/model...")
                continue
            except Exception as e:
                last_error = f"Error with {model_name} on {api_version}: {e}"
                st.warning(f"Debug: {last_error} Trying fallback endpoint/model...")
                continue

    st.error(f"All Gemini models failed or are currently experiencing high demand. Last error: {last_error}")
    return None, None, None

# --- Streamlit UI ---
st.set_page_config(page_title="Junior Editor", layout="wide")
st.title("Junior Editor")

st.markdown("""
**Instructions**
* Upload your file here (video or audio).
* Set your timeline FPS, Source Timecode, and transcription quality in the sidebar.
* Junior Editor will transcribe and separate speakers. You can then instruct it to find engaging bits or construct a narrative.
* **Note: If you want me to write voice over suggestions or title cards, please request it.**
* It will create an EDL or XML to import back into your editing software (Resolve, Premiere).
* **For Multicam Workflows:** Use the **XML** option in the sidebar and enter the **exact name** of your Multicam Sequence.
* **Resolve Users:** Uncheck "Automatically import source clips into media pool" during import.
* **Premiere Users:** Use the EDL option if XML conformance fails for Multicam sequences.
""")

st.divider()

with st.sidebar:
    clip_settings_container = st.container()
    transcription_settings_container = st.container()

    with clip_settings_container:
        st.header("Clip Settings")
        source_start_tc = st.text_input("Source Start Timecode", value="00:00:00:00", help="Format: HH:MM:SS:FF")
        
        export_format = st.radio("Output Format", ["EDL", "XML (Multicam)"], index=0)
        
        input_label = "EDL Reel Name"
        input_help = "Leave empty to use the uploaded file name."
        if export_format == "XML (Multicam)":
            input_label = "Multicam Sequence Name"
            input_help = "EXACT name of your Multicam Clip in Resolve."
        
        custom_reel_name = st.text_input(
            input_label, 
            placeholder="e.g. Interview_Day1_Multi", 
            help=input_help
        )
        
        fps_options = [23.98, 24, 25, 29.97, 30, 50, 59.94, 60]
        fps = st.selectbox("Timeline FPS", fps_options, index=2)

    with transcription_settings_container:
        st.header("Transcription Settings")
        
        language_map = {
            "Auto-Detect": None,
            "English": "en",
            "Spanish": "es",
            "French": "fr",
            "German": "de",
            "Italian": "it",
            "Portuguese": "pt"
        }
        selected_lang_label = st.selectbox("Audio Language", list(language_map.keys()), index=1)
        target_language = language_map[selected_lang_label]
        
        num_speakers = st.number_input("Speakers (0=Auto)", min_value=0, value=0)

uploaded_file = st.file_uploader("Upload Video/Audio Clip", type=["mp4", "m4a", "wav", "mp3", "mov"])

if uploaded_file:
    # --- Auto-Reset Logic ---
    if "last_processed_file" not in st.session_state or st.session_state.last_processed_file != uploaded_file.name:
        if "transcript" in st.session_state:
            del st.session_state.transcript
        if "edit_generated" in st.session_state:
            del st.session_state.edit_generated
        st.session_state.last_processed_file = uploaded_file.name

    # --- Auto-Process Logic ---
    if "transcript" not in st.session_state:
        if not ACTIVE_HF_TOKEN or "PASTE_YOUR_HF_TOKEN" in ACTIVE_HF_TOKEN:
            st.error("Please provide a valid Hugging Face Token in the Sidebar/Secrets.")
        else:
            progress_container = st.container()
            with progress_container:
                st.info("🤖 **Junior Editor is processing your file...**")
                status_text = st.empty()
                progress_bar = st.progress(0)
                
                try:
                    # Phase 1
                    status_text.markdown("**Phase 1/4: Extracting Audio...**")
                    with open("temp_input", "wb") as f:
                        f.write(uploaded_file.getbuffer())
                    
                    subprocess.run(["ffmpeg", "-i", "temp_input", "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1", "temp_audio.wav", "-y"])
                    progress_bar.progress(25)
                    
                    device = "cuda" if torch.cuda.is_available() else "cpu"
                    if device == "cpu": st.warning("⚠️ No GPU detected.")
                    
                    # Phase 2
                    status_text.markdown("**Phase 2/4: Transcribing (Whisper)... This is the longest step.**")
                    compute_type = "float16" if device == "cuda" else "int8"
                    model = whisperx.load_model("large-v2", device, compute_type=compute_type)
                    audio = whisperx.load_audio("temp_audio.wav")
                    result = model.transcribe(audio, batch_size=4, language=target_language)
                    del model
                    gc.collect()
                    torch.cuda.empty_cache()
                    progress_bar.progress(50)
                    
                    # Phase 3
                    status_text.markdown("**Phase 3/4: Aligning Text...**")
                    model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=device)
                    result = whisperx.align(result["segments"], model_a, metadata, audio, device, return_char_alignments=False)
                    del model_a
                    gc.collect()
                    torch.cuda.empty_cache()
                    progress_bar.progress(75)
                    
                    # Phase 4
                    status_text.markdown("**Phase 4/4: Identifying Speakers...**")
                    diarize_model = whisperx.DiarizationPipeline(use_auth_token=ACTIVE_HF_TOKEN, device=device)
                    diarize_kwargs = {"min_speakers": num_speakers, "max_speakers": num_speakers} if num_speakers > 0 else {}
                    diarize_segments = diarize_model(audio, **diarize_kwargs)
                    
                    # Final Merge
                    status_text.markdown("**Finalizing...**")
                    final_result = whisperx.assign_word_speakers(diarize_segments, result)
                    
                    processed_segments = []
                    for segment in final_result["segments"]:
                        processed_segments.append({
                            "speaker": segment.get("speaker", "Unknown"),
                            "text": segment["text"].strip(),
                            "start": segment["start"],
                            "end": segment["end"]
                        })
                    
                    st.session_state.transcript = processed_segments
                    if os.path.exists("temp_input"): os.remove("temp_input")
                    if os.path.exists("temp_audio.wav"): os.remove("temp_audio.wav")
                    progress_bar.progress(100)
                    status_text.success(f"Done! Processed {len(processed_segments)} segments.")

                except Exception as e:
                    status_text.error(f"Error: {e}")
                    if os.path.exists("temp_input"): os.remove("temp_input")
                    st.stop()

    if "transcript" in st.session_state:
        st.divider()
        with st.expander("Transcript Preview", expanded=False):
            for seg in st.session_state.transcript: 
                st.markdown(f"**{seg['speaker']}:** {seg['text']}")
        
        st.subheader("Your Instruction")
        brief = st.text_area("What should the Junior Editor do?", placeholder="e.g. Find the most engaging bits. If you want me to write voice over suggestions or title cards, please request it.")
        
        if st.button("Generate Edit"):
            if not ACTIVE_GEMINI_KEY:
                st.error("Gemini API Key required.")
            else:
                with st.spinner("Junior Editor is thinking..."):
                    final_source_name = custom_reel_name.strip() if custom_reel_name.strip() else uploaded_file.name
                    offset_frames = timecode_to_frames(source_start_tc, fps)
                    
                    edl_segments, used_model, used_endpoint = call_gemini_for_edl(st.session_state.transcript, brief, ACTIVE_GEMINI_KEY)
                    
                    if edl_segments:
                        safe_filename = "".join([c for c in final_source_name if c.isalnum() or c in (' ', '_', '-')]).strip()
                        safe_filename = safe_filename.replace(' ', '_')
                        if not safe_filename:
                            safe_filename = "junior_editor_cut"
                            
                        # Generate the selected edit format (EDL/XML)
                        if export_format == "EDL":
                            final_output = generate_cmx_edl(final_source_name, edl_segments, final_source_name, fps, offset_frames)
                            ext = "edl"
                        else:
                            final_output = generate_xml(final_source_name, edl_segments, final_source_name, fps, offset_frames)
                            ext = "xml"
                            
                        # Generate the companion human-readable transcript
                        transcript_txt = generate_transcript_txt(final_source_name, edl_segments, fps, offset_frames, st.session_state.transcript)
                        
                        # Save to session state so buttons don't disappear on click
                        st.session_state.edit_generated = True
                        st.session_state.revision_count = 1
                        st.session_state.base_filename = safe_filename
                        st.session_state.final_output = final_output
                        st.session_state.ext = ext
                        st.session_state.transcript_txt = transcript_txt
                        st.session_state.safe_filename = safe_filename
                        st.session_state.used_model = used_model
                        st.session_state.used_endpoint = used_endpoint
                        st.session_state.edl_segments = edl_segments
                        
        # Render the download buttons outside the Generate block using session state
        if st.session_state.get("edit_generated"):
            st.subheader("Ready for Import")
            st.success(f"✅ Edit generated successfully using **{st.session_state.used_model}** (via {st.session_state.used_endpoint}).")
            
            # Display downloads in a neat row
            col1, col2 = st.columns(2)
            with col1:
                st.download_button(f"Download .{st.session_state.ext.upper()} Sequence", data=st.session_state.final_output, file_name=f"{st.session_state.safe_filename}.{st.session_state.ext}")
            with col2:
                st.download_button(f"Download Reference Transcript (.TXT)", data=st.session_state.transcript_txt, file_name=f"{st.session_state.safe_filename}-Transcript.txt")
            
            st.info("📝 **Note:** There seems to be a bug in Resolve that the first time you import the sequence into a bin, sometimes it won't relink. Just Delete the sequence and Import again and it should work.")
            
            st.divider()
            st.subheader("Revise Edit")
            revision_brief = st.text_area("Want changes? Tell Junior Editor:", placeholder="e.g. Make it twice as long, add a graphic card for location, or write a VO intro.")
            
            if st.button("Apply Revisions"):
                with st.spinner("Junior Editor is revising the edit..."):
                    final_source_name = custom_reel_name.strip() if custom_reel_name.strip() else uploaded_file.name
                    offset_frames = timecode_to_frames(source_start_tc, fps)
                    
                    new_edl_segments, used_model, used_endpoint = call_gemini_for_edl(
                        st.session_state.transcript, 
                        revision_brief, 
                        ACTIVE_GEMINI_KEY, 
                        previous_edit=st.session_state.edl_segments
                    )
                    
                    if new_edl_segments:
                        # Increment version count
                        st.session_state.revision_count = st.session_state.get("revision_count", 1) + 1
                        rev_seq_name = f"{final_source_name} V{st.session_state.revision_count}"
                        
                        if export_format == "EDL":
                            final_output = generate_cmx_edl(rev_seq_name, new_edl_segments, final_source_name, fps, offset_frames)
                            ext = "edl"
                        else:
                            final_output = generate_xml(rev_seq_name, new_edl_segments, final_source_name, fps, offset_frames)
                            ext = "xml"
                            
                        transcript_txt = generate_transcript_txt(rev_seq_name, new_edl_segments, fps, offset_frames, st.session_state.transcript)
                        
                        st.session_state.final_output = final_output
                        st.session_state.ext = ext
                        st.session_state.transcript_txt = transcript_txt
                        st.session_state.used_model = used_model
                        st.session_state.used_endpoint = used_endpoint
                        st.session_state.edl_segments = new_edl_segments
                        
                        # Update safe filename for the download buttons to show _V2, _V3, etc.
                        st.session_state.safe_filename = f"{st.session_state.base_filename}_V{st.session_state.revision_count}"
                        
                        st.rerun()