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 re
from urllib.parse import quote

# --- 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):
    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):
    if not text: return ""
    text = str(text)
    return text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;").replace('"', "&quot;").replace("'", "&apos;")

def format_tc_string(tc):
    """Automatically adds colons to 8-digit timecodes (e.g., 01000000 -> 01:00:00:00)."""
    tc_clean = re.sub(r'[^\d]', '', tc)
    if len(tc_clean) >= 8:
        return f"{tc_clean[0:2]}:{tc_clean[2:4]}:{tc_clean[4:6]}:{tc_clean[6:8]}"
    return tc

def timecode_to_frames(tc, fps):
    tc = format_tc_string(tc)
    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)) 
    return (h * 3600 * timebase) + (m * 60 * timebase) + (s * timebase) + f

def frames_to_timecode(frames, fps):
    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):
    return int(round(seconds * fps))

# --- MULTI-CLIP EXPORT GENERATORS ---

def generate_cmx_edl(edl_title, segments, clip_metadata, fps=25):
    """Constructs a CMX 3600 formatted EDL handling multiple source clips."""
    edl_lines = [f"TITLE: {edl_title}", "FCM: NON-DROP FRAME\n"]
    rec_start_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
            current_source = "GEN_VO" if is_vo else "GEN_GFX"
            note_text = f"{'VO SCRIPT' if is_vo else 'GRAPHIC'}: {text_content}"
            source_name = "Generated"
        else:
            source_clip = seg.get('source_clip', 'UNKNOWN')
            meta = clip_metadata.get(source_clip, {'offset_frames': 0})
            start_offset_frames = meta['offset_frames']
            
            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 = source_clip.replace(" ", "_")[:8]
            note_text = seg.get('note', 'Clip')
            source_name = source_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, clip_metadata, fps=25):
    """Constructs a Robust XML referencing multiple source files dynamically."""
    timebase = int(round(fps))
    is_ntsc = "TRUE" if fps % 1 != 0 else "FALSE"

    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>')

    defined_files = set()
    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
            
            source_clip = "Generated_VO_Placeholder" if is_vo else "Generated_Graphic_Placeholder"
            master_id_to_use = f"masterfile-{source_clip.lower().replace('_', '')}"
            start_offset_frames = 0
            start_tc_string = "00:00:00:00"
            path_url = f"file://localhost/{escape_xml(source_clip)}.{'wav' if is_vo else 'mov'}"
            
            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:
            source_clip = seg.get('source_clip', 'UNKNOWN')
            meta = clip_metadata.get(source_clip, {'offset_frames': 0, 'start_tc': '00:00:00:00'})
            start_offset_frames = meta['offset_frames']
            start_tc_string = meta['start_tc']
            
            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
            
            safe_id = "".join([c for c in source_clip if c.isalnum()]).lower()
            master_id_to_use = f"masterfile-{safe_id}"
            path_url = f"file://localhost/{quote(source_clip)}"
            
            clip_note = seg.get('note', 'Junior Editor Selection')
            clip_name_disp = source_clip
        
        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>')
        
        if master_id_to_use not in defined_files:
            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_clip)}</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>')
            if not is_graphic: 
                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>')
            defined_files.add(master_id_to_use)
        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>')
        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:
            duration_sec = float(seg.get('duration', 4.0))
            timeline_head_frames += seconds_to_frames(duration_sec, fps) + seconds_to_frames(float(seg.get('gap', 0.0)), fps)
            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
            source_clip = "Generated_VO_Placeholder"
        else:
            source_clip = seg.get('source_clip', 'UNKNOWN')
            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
            
        safe_id = "".join([c for c in source_clip if c.isalnum()]).lower()
        master_id_to_use = f"masterfile-{safe_id}"
            
        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_clip)}</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, clip_metadata, fps=25, all_transcripts_data=None):
    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:
            source_clip = seg.get('source_clip', 'UNKNOWN')
            meta = clip_metadata.get(source_clip, {'offset_frames': 0})
            
            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) + meta['offset_frames']
            src_out_frames = seconds_to_frames(src_end_val, fps) + meta['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 all_transcripts_data and source_clip in all_transcripts_data:
                clip_words = []
                for t_seg in all_transcripts_data[source_clip]:
                    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}  |  SOURCE: {source_clip}  |  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(transcripts_dict, story_prompt, api_key, previous_edit=None):
    if not api_key:
        st.error("Gemini API Key is missing.")
        return None, None, None
        
    system_prompt = (
        "You are an expert Documentary Senior Editor. Use the provided JSON object containing transcripts "
        "from ONE OR MORE source clips to create a condensed story.\n\n"
        "STRICT INSTRUCTIONS ON GENERATIVE ELEMENTS:\n"
        "- DO NOT generate 'vo' (Voice-Over) or 'graphic' (Title Cards) segments UNLESS the user explicitly asks for them in their creative brief.\n"
        "- By default, you must ONLY construct the edit using actual extracted 'clip' segments from the provided transcripts.\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\", \"source_clip\": \"Cam_A_Interview\", \"src_start\": 12.5, \"src_end\": 25.0, \"note\": \"Subject talks about X\", \"gap\": 0.0}\n"
        "- MULTI-CAM CRITICAL: You MUST include the exact 'source_clip' key representing which file the transcript came from.\n"
        "- IGNORE ALL INTERVIEWER COMMENTS.\n"
        "- REMOVE FLUFF: Delete 'um', 'ah', repeats.\n"
        "- TIMESTAMP INTEGRITY: Use exact word-level start/end times.\n\n"
        "For 'vo' segments (ONLY IF EXPLICITLY REQUESTED):\n"
        "{\"type\": \"vo\", \"text\": \"The journey began years ago...\", \"duration\": 3.5, \"gap\": 0.0}\n\n"
        "For 'graphic' segments (ONLY IF EXPLICITLY REQUESTED):\n"
        "{\"type\": \"graphic\", \"text\": \"Chapter 1\", \"duration\": 4.0, \"gap\": 0.0}\n\n"
        "PACING: Add a 'gap' (in seconds) between distinct ideas."
    )
    
    if previous_edit:
        prompt_text = (
            f"Transcripts Data (By Source):\n{json.dumps(transcripts_dict)}\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 applying these requested changes."
        )
    else:
        prompt_text = f"Creative Brief: {story_prompt}\n\nTranscripts Data (By Source):\n{json.dumps(transcripts_dict)}"
    
    payload = {
        "contents": [{"role": "user", "parts": [{"text": prompt_text}]}],
        "systemInstruction": {"parts": [{"text": system_prompt}]},
        "generationConfig": {"responseMimeType": "application/json"}
    }
    
    headers = {"Content-Type": "application/json"}
    
    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)
                if res.status_code == 200:
                    result_json = res.json()
                    raw_text = result_json['candidates'][0]['content']['parts'][0]['text']
                    return json.loads(clean_json_response(raw_text)), model_name, api_version
                elif res.status_code in [503, 429]:
                    last_error = f"Model {model_name} busy ({res.status_code}) on {api_version}."
                    continue 
                else:
                    last_error = f"API rejected ({res.status_code}) on {api_version}."
                    continue 
            except Exception as e:
                last_error = f"Error with {model_name} on {api_version}: {e}"
                continue

    st.error(f"All Gemini models failed. Last error: {last_error}")
    return None, None, None

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

with st.expander("Instructions & Prompt Guide - Please Read First", expanded=False):
    st.markdown("""
    * This program will generate an EDL or XML containing multiple interview cut down and interwoven as directed by you.
    * Set your export timeline FPS and EDL or XML. If you are doing sequences it will have to be XML, but this will only work in Resolve (I haven’t tried FCP). If you are exporting sequences have the TC start at 00:00:00:00 for speed later. 
    * First, export compressed source files from your timeline. These can be sequences (inc multicam) or individual clips.
    * Drag and drop multiple compressed files below (video or audio). AAC MP4’s work well for speed of upload. 
    * Enter the exact name of each your sequence or clip as well as the start timecode of the reference. You can add TC’s as `01000000` and the program will add colons for you. 
    * Junior Editor will transcribe and separate speakers. This takes approx 10 mins per hour of audio. Once completed.
    * You can then instruct it to find engaging bits or construct a narrative. You can ask it to add in suggested VO or Graphics cards if you are planning to use voice over. 
    * It will create an EDL or XML to import back into Resolve. 
    * Resolve Users: Uncheck "Automatically import source clips into media pool" during import.
    
    ---
    ### 💡 Pro-Tips for Prompting the AI
    **1. Think in Clips, Not Minutes:** The AI cannot accurately calculate running time. If you want a video longer than 3-5 minutes, **do not ask for a time limit**. Instead, ask for a specific number of clips (assume roughly 8-10 clips per minute).
    
    **2. Use Bullet Points:** The AI follows technical, structured lists much better than paragraphs.
    
    **3. Do the Math for It:** If you want specific ratios, give the AI hard numbers (e.g., "Use 30 clips from Cam A and 10 from Cam B") instead of percentages (e.g., "75% Cam A").
    
    **Example of a Great Prompt:**
    > **TASK:** Build a comprehensive highlights reel from the provided event transcripts. 
    > 
    > **LENGTH CONSTRAINT:** You MUST generate a minimum of 75 to 80 individual clip segments. Do not stop early.
    > 
    > **SOURCE ALLOCATION:**
    > * **~32 clips** from `BeutlerClips2` (Use heavily as the inspiring backbone).
    > * **~8 clips** from `Bruce Interview`.
    > * **~40 clips** mixed from the other attendee interviews. 
    > 
    > **NARRATIVE & PACING RULES:**
    > * **No VO/Graphics.** > * **The Hook:** Start with an intriguing hook, followed by a clip explaining the event.
    > * **The Chain:** Each clip must organically lead into the next. Do not repeat the same info. 
    > * **The Ending:** Finish with a highly uplifting statement.
    """)

st.divider()

# -- Global Settings Sidebar --
with st.sidebar:
    st.header("Project Settings")
    sequence_title = st.text_input("Sequence / Project Title", value="Junior_Editor_Cut")
    export_format = st.radio("Output Format", ["XML (Premiere/Resolve)", "EDL"], index=0)
    
    fps_options = [23.98, 24, 25, 29.97, 30, 50, 59.94, 60]
    fps = st.selectbox("Timeline FPS", fps_options, index=2)
    
    st.header("Transcription Settings")
    language_map = {"Auto-Detect": None, "English": "en", "Spanish": "es"}
    target_language = language_map[st.selectbox("Audio Language", list(language_map.keys()), index=1)]

# -- Main Multi-Clip Bin area --
if "transcripts" not in st.session_state:
    st.session_state.transcripts = {}
if "clip_metadata" not in st.session_state:
    st.session_state.clip_metadata = {}
if "revision_count" not in st.session_state:
    st.session_state.revision_count = 1

uploaded_files = st.file_uploader("Upload Clip(s)", type=["mp4", "m4a", "wav", "mp3", "mov"], accept_multiple_files=True)

if uploaded_files and len(st.session_state.transcripts) == 0:
    st.subheader("Clip Settings Bin")
    clip_configs = {}
    
    for idx, f in enumerate(uploaded_files):
        with st.expander(f"⚙️ Settings: {f.name}", expanded=True):
            col1, col2 = st.columns(2)
            with col1:
                custom_name = st.text_input(
                    "Clip / Multicam Name", 
                    value=f.name.rsplit('.', 1)[0], 
                    key=f"name_{idx}",
                    help="The EXACT name of the clip or Multicam sequence as it appears in your NLE Bin."
                )
            with col2:
                start_tc = st.text_input(
                    "Source Start Timecode", 
                    value="00:00:00:00", 
                    key=f"tc_{idx}",
                    help="Format HH:MM:SS:FF or HHMMSSFF. We will auto-format it."
                )
            clip_configs[f.name] = {"file": f, "custom_name": custom_name, "start_tc": start_tc}

    if st.button("Transcribe All Clips", type="primary"):
        if not ACTIVE_HF_TOKEN or "PASTE_YOUR_HF_TOKEN" in ACTIVE_HF_TOKEN:
            st.error("Please provide a valid Hugging Face Token.")
            st.stop()
            
        progress_bar = st.progress(0)
        status_text = st.empty()
        
        all_transcripts = {}
        all_metadata = {}
        
        device = "cuda" if torch.cuda.is_available() else "cpu"
        compute_type = "float16" if device == "cuda" else "int8"
        
        status_text.markdown("**Loading AI Models into memory...**")
        whisper_model = whisperx.load_model("large-v2", device, compute_type=compute_type)
        diarize_model = whisperx.DiarizationPipeline(use_auth_token=ACTIVE_HF_TOKEN, device=device)
        
        total_clips = len(clip_configs)
        
        for i, (original_filename, config) in enumerate(clip_configs.items()):
            c_name = config["custom_name"]
            raw_tc = config["start_tc"]
            c_file = config["file"]
            
            status_text.markdown(f"**Processing Clip {i+1}/{total_clips}: {c_name}...**")
            
            with open("temp_input", "wb") as f: f.write(c_file.getbuffer())
            subprocess.run(["ffmpeg", "-i", "temp_input", "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1", "temp_audio.wav", "-y"], capture_output=True)
            
            audio = whisperx.load_audio("temp_audio.wav")
            result = whisper_model.transcribe(audio, batch_size=4, language=target_language)
            
            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()
            
            diarize_segments = diarize_model(audio) 
            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"]
                })
            
            formatted_tc = format_tc_string(raw_tc)
            
            all_transcripts[c_name] = processed_segments
            all_metadata[c_name] = {
                "start_tc": formatted_tc, 
                "offset_frames": timecode_to_frames(formatted_tc, fps)
            }
            
            progress_bar.progress((i + 1) / total_clips)
            
        # Cleanup
        del whisper_model
        del diarize_model
        gc.collect()
        torch.cuda.empty_cache()
        if os.path.exists("temp_input"): os.remove("temp_input")
        if os.path.exists("temp_audio.wav"): os.remove("temp_audio.wav")
        
        st.session_state.transcripts = all_transcripts
        st.session_state.clip_metadata = all_metadata
        status_text.success("All clips transcribed successfully!")
        st.rerun()

# -- Chat / Generation Area --
if len(st.session_state.transcripts) > 0 and not st.session_state.get("edit_generated"):
    st.divider()
    
    with st.expander("View Source Transcripts", expanded=False):
        for clip_name, segments in st.session_state.transcripts.items():
            st.markdown(f"### {clip_name}")
            for seg in segments[:10]: # Preview first 10
                st.markdown(f"**{seg['speaker']}**: {seg['text']}")
            if len(segments) > 10: st.markdown("*... (truncated)*")
            st.divider()

    st.subheader("Your Instruction")
    brief = st.text_area("What should the Junior Editor do?", placeholder="e.g. Build a sequence using 40 clips, starting with an intriguing hook...")
    
    if st.button("Generate Multi-Clip Edit", type="primary"):
        with st.spinner("Analyzing all transcripts and building sequence..."):
            edl_segments, used_model, used_endpoint = call_gemini_for_edl(st.session_state.transcripts, brief, ACTIVE_GEMINI_KEY)
            
            if edl_segments:
                safe_title = "".join([c for c in sequence_title if c.isalnum() or c in (' ', '_', '-')]).strip().replace(' ', '_')
                
                if "XML" in export_format:
                    final_output = generate_xml(f"{safe_title}_V1", edl_segments, st.session_state.clip_metadata, fps)
                    ext = "xml"
                else:
                    final_output = generate_cmx_edl(f"{safe_title}_V1", edl_segments, st.session_state.clip_metadata, fps)
                    ext = "edl"
                    
                transcript_txt = generate_transcript_txt(f"{safe_title}_V1", edl_segments, st.session_state.clip_metadata, fps, st.session_state.transcripts)
                
                st.session_state.edit_generated = True
                st.session_state.revision_count = 1
                st.session_state.safe_title = safe_title
                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.edl_segments = edl_segments
                st.rerun()

# -- Post-Generation / Revision Area --
if st.session_state.get("edit_generated"):
    st.success(f"✅ Multi-clip sequence generated successfully using **{st.session_state.used_model}**.")
    
    col1, col2 = st.columns(2)
    with col1:
        st.download_button(f"Download .{st.session_state.ext.upper()}", data=st.session_state.final_output, file_name=f"{st.session_state.safe_title}_V{st.session_state.revision_count}.{st.session_state.ext}")
    with col2:
        st.download_button(f"Download Transcript (.TXT)", data=st.session_state.transcript_txt, file_name=f"{st.session_state.safe_title}_V{st.session_state.revision_count}_Transcript.txt")

    st.divider()
    
    st.subheader("Revise This Edit")
    revision_brief = st.text_area("Want changes? Give Junior Editor new instructions based on the current cut:", placeholder="e.g. Make it longer (add 20 more clips), swap the opening clip, or focus more on AstraZeneca leaders.")
    
    col3, col4 = st.columns([1, 4])
    with col3:
        if st.button("Apply Revisions", type="primary"):
            with st.spinner("Junior Editor is revising the sequence..."):
                new_edl_segments, used_model, used_endpoint = call_gemini_for_edl(
                    st.session_state.transcripts, 
                    revision_brief, 
                    ACTIVE_GEMINI_KEY, 
                    previous_edit=st.session_state.edl_segments
                )
                
                if new_edl_segments:
                    st.session_state.revision_count += 1
                    safe_title = st.session_state.safe_title
                    
                    if st.session_state.ext == "xml":
                        final_output = generate_xml(f"{safe_title}_V{st.session_state.revision_count}", new_edl_segments, st.session_state.clip_metadata, fps)
                    else:
                        final_output = generate_cmx_edl(f"{safe_title}_V{st.session_state.revision_count}", new_edl_segments, st.session_state.clip_metadata, fps)
                        
                    transcript_txt = generate_transcript_txt(f"{safe_title}_V{st.session_state.revision_count}", new_edl_segments, st.session_state.clip_metadata, fps, st.session_state.transcripts)
                    
                    st.session_state.final_output = final_output
                    st.session_state.transcript_txt = transcript_txt
                    st.session_state.used_model = used_model
                    st.session_state.edl_segments = new_edl_segments
                    st.rerun()
    with col4:
        if st.button("Start Over / Clear Bin"):
            for key in ["transcripts", "clip_metadata", "edit_generated", "revision_count"]:
                if key in st.session_state: del st.session_state[key]
            st.rerun()