import streamlit as st
import anthropic
import openai
import base64
import cv2
import glob
import json
import math
import os
import pytz
import random
import re
import requests
import textract
import time
import zipfile
import plotly.graph_objects as go
import streamlit.components.v1 as components
from datetime import datetime
from audio_recorder_streamlit import audio_recorder
from bs4 import BeautifulSoup
from collections import defaultdict, deque
from dotenv import load_dotenv
from gradio_client import Client
from huggingface_hub import InferenceClient
from io import BytesIO
from PIL import Image
from PyPDF2 import PdfReader
from urllib.parse import quote
from xml.etree import ElementTree as ET
from openai import OpenAI
import extra_streamlit_components as stx
from streamlit.runtime.scriptrunner import get_script_run_ctx
import asyncio
import edge_tts

# 🎯 1. Core Configuration & Setup
st.set_page_config(
    page_title="🚲BikeAI🏆 Claude/GPT Research",
    page_icon="🚲🏆",
    layout="wide",
    initial_sidebar_state="auto",
    menu_items={
        'Get Help': 'https://huggingface.co/awacke1',
        'Report a bug': 'https://huggingface.co/spaces/awacke1',
        'About': "🚲BikeAI🏆 Claude/GPT Research AI"
    }
)
load_dotenv()

# 🔑 2. API Setup & Clients
openai_api_key = os.getenv('OPENAI_API_KEY', "")
anthropic_key = os.getenv('ANTHROPIC_API_KEY_3', "")
xai_key = os.getenv('xai',"")
if 'OPENAI_API_KEY' in st.secrets:
    openai_api_key = st.secrets['OPENAI_API_KEY']
if 'ANTHROPIC_API_KEY' in st.secrets:
    anthropic_key = st.secrets["ANTHROPIC_API_KEY"]

openai.api_key = openai_api_key
claude_client = anthropic.Anthropic(api_key=anthropic_key)
openai_client = OpenAI(api_key=openai.api_key, organization=os.getenv('OPENAI_ORG_ID'))
HF_KEY = os.getenv('HF_KEY')
API_URL = os.getenv('API_URL')

# 📝 3. Session State Management
if 'transcript_history' not in st.session_state:
    st.session_state['transcript_history'] = []
if 'chat_history' not in st.session_state:
    st.session_state['chat_history'] = []
if 'openai_model' not in st.session_state:
    st.session_state['openai_model'] = "gpt-4o-2024-05-13"
if 'messages' not in st.session_state:
    st.session_state['messages'] = []
if 'last_voice_input' not in st.session_state:
    st.session_state['last_voice_input'] = ""
if 'editing_file' not in st.session_state:
    st.session_state['editing_file'] = None
if 'edit_new_name' not in st.session_state:
    st.session_state['edit_new_name'] = ""
if 'edit_new_content' not in st.session_state:
    st.session_state['edit_new_content'] = ""
if 'viewing_prefix' not in st.session_state:
    st.session_state['viewing_prefix'] = None
if 'should_rerun' not in st.session_state:
    st.session_state['should_rerun'] = False
if 'old_val' not in st.session_state:
    st.session_state['old_val'] = None

# 🎨 4. Custom CSS
st.markdown("""
<style>
    .main { background: linear-gradient(to right, #1a1a1a, #2d2d2d); color: #fff; }
    .stMarkdown { font-family: 'Helvetica Neue', sans-serif; }
    .stButton>button {
        margin-right: 0.5rem;
    }
</style>
""", unsafe_allow_html=True)

FILE_EMOJIS = {
    "md": "📝",
    "mp3": "🎵",
}

# 🧠 5. High-Information Content Extraction
def get_high_info_terms(text: str) -> list:
    """Extract high-information terms from text, including key phrases."""
    stop_words = set([
        'the', 'a', 'an', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for', 'of', 'with',
        'by', 'from', 'up', 'about', 'into', 'over', 'after', 'is', 'are', 'was', 'were',
        'be', 'been', 'being', 'have', 'has', 'had', 'do', 'does', 'did', 'will', 'would',
        'should', 'could', 'might', 'must', 'shall', 'can', 'may', 'this', 'that', 'these',
        'those', 'i', 'you', 'he', 'she', 'it', 'we', 'they', 'what', 'which', 'who',
        'when', 'where', 'why', 'how', 'all', 'any', 'both', 'each', 'few', 'more', 'most',
        'other', 'some', 'such', 'than', 'too', 'very', 'just', 'there'
    ])

    key_phrases = [
        'artificial intelligence', 'machine learning', 'deep learning', 'neural network',
        'personal assistant', 'natural language', 'computer vision', 'data science',
        'reinforcement learning', 'knowledge graph', 'semantic search', 'time series',
        'large language model', 'transformer model', 'attention mechanism',
        'autonomous system', 'edge computing', 'quantum computing', 'blockchain technology',
        'cognitive science', 'human computer', 'decision making', 'arxiv search',
        'research paper', 'scientific study', 'empirical analysis'
    ]

    # Identify key phrases
    preserved_phrases = []
    lower_text = text.lower()
    for phrase in key_phrases:
        if phrase in lower_text:
            preserved_phrases.append(phrase)
            text = text.replace(phrase, '')

    # Extract individual words
    words = re.findall(r'\b\w+(?:-\w+)*\b', text)
    high_info_words = [
        word.lower() for word in words 
        if len(word) > 3
        and word.lower() not in stop_words
        and not word.isdigit()
        and any(c.isalpha() for c in word)
    ]

    all_terms = preserved_phrases + high_info_words
    seen = set()
    unique_terms = []
    for term in all_terms:
        if term not in seen:
            seen.add(term)
            unique_terms.append(term)

    max_terms = 5
    return unique_terms[:max_terms]

def clean_text_for_filename(text: str) -> str:
    """Remove punctuation and short filler words, return a compact string."""
    text = text.lower()
    text = re.sub(r'[^\w\s-]', '', text)
    words = text.split()
    stop_short = set(['the','and','for','with','this','that','from','just','very','then','been','only','also','about'])
    filtered = [w for w in words if len(w)>3 and w not in stop_short]
    return '_'.join(filtered)[:200]

# 📁 6. File Operations
def generate_filename(prompt, response, file_type="md"):
    """
    Generate filename with meaningful terms and short dense clips from prompt & response.
    The filename should be about 150 chars total, include high-info terms, and a clipped snippet.
    """
    prefix = datetime.now().strftime("%y%m_%H%M") + "_"
    combined = (prompt + " " + response).strip()
    info_terms = get_high_info_terms(combined)
    
    # Include a short snippet from prompt and response
    snippet = (prompt[:100] + " " + response[:100]).strip()
    snippet_cleaned = clean_text_for_filename(snippet)
    
    # Combine info terms and snippet
    name_parts = info_terms + [snippet_cleaned]
    full_name = '_'.join(name_parts)

    # Trim to ~150 chars
    if len(full_name) > 150:
        full_name = full_name[:150]
    
    filename = f"{prefix}{full_name}.{file_type}"
    return filename

def create_file(prompt, response, file_type="md"):
    """Create file with an intelligent naming scheme."""
    filename = generate_filename(prompt.strip(), response.strip(), file_type)
    with open(filename, 'w', encoding='utf-8') as f:
        f.write(prompt + "\n\n" + response)
    return filename

def get_download_link(file):
    """Generate download link for file"""
    with open(file, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()
    return f'<a href="data:file/zip;base64,{b64}" download="{os.path.basename(file)}">📂 Download {os.path.basename(file)}</a>'

# 🔊 7. Audio Processing
def clean_for_speech(text: str) -> str:
    """Clean text for speech synthesis"""
    text = text.replace("\n", " ")
    text = text.replace("</s>", " ")
    text = text.replace("#", "")
    text = re.sub(r"\(https?:\/\/[^\)]+\)", "", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text

@st.cache_resource
def speech_synthesis_html(result):
    """Create HTML for speech synthesis"""
    html_code = f"""
    <html><body>
    <script>
    var msg = new SpeechSynthesisUtterance("{result.replace('"', '')}");
    window.speechSynthesis.speak(msg);
    </script>
    </body></html>
    """
    components.html(html_code, height=0)

async def edge_tts_generate_audio(text, voice="en-US-AriaNeural", rate=0, pitch=0):
    """Generate audio using Edge TTS (async)"""
    text = clean_for_speech(text)
    if not text.strip():
        return None
    rate_str = f"{rate:+d}%"
    pitch_str = f"{pitch:+d}Hz"
    communicate = edge_tts.Communicate(text, voice, rate=rate_str, pitch=pitch_str)
    out_fn = generate_filename(text, text, "mp3")
    await communicate.save(out_fn)
    return out_fn

def speak_with_edge_tts(text, voice="en-US-AriaNeural", rate=0, pitch=0):
    """Wrapper for edge TTS generation (sync)"""
    return asyncio.run(edge_tts_generate_audio(text, voice, rate, pitch))

def play_and_download_audio(file_path):
    """Play and provide a download link for audio"""
    if file_path and os.path.exists(file_path):
        st.audio(file_path)
        dl_link = f'<a href="data:audio/mpeg;base64,{base64.b64encode(open(file_path,"rb").read()).decode()}" download="{os.path.basename(file_path)}">Download {os.path.basename(file_path)}</a>'
        st.markdown(dl_link, unsafe_allow_html=True)

def auto_play_audio(file_path):
    """
    Reads MP3 file as base64, displays an <audio> tag with autoplay + controls + download link.
    Note: Some browsers block audio autoplay if there's no user interaction.
    """
    if not file_path or not os.path.exists(file_path):
        return
    with open(file_path, "rb") as f:
        b64_data = base64.b64encode(f.read()).decode("utf-8")
    filename = os.path.basename(file_path)
    st.markdown(f"""
    <audio controls autoplay>
        <source src="data:audio/mpeg;base64,{b64_data}" type="audio/mpeg">
        Your browser does not support the audio element.
    </audio>
    <br/>
    <a href="data:audio/mpeg;base64,{b64_data}" download="{filename}">
        Download {filename}
    </a>
    """, unsafe_allow_html=True)

def generate_audio_filename(query, title, summary):
    """
    Example specialized MP3 filename: prefix + query + short snippet of title/summary
    """
    combined = (query + " " + title + " " + summary).strip().lower()
    combined = re.sub(r'[^\w\s-]', '', combined)  # remove special chars
    combined = "_".join(combined.split())[:80]    # limit length
    prefix = datetime.now().strftime("%y%m_%H%M")
    return f"{prefix}_{combined}.mp3"

# 🎬 8. Media Processing
def process_image(image_path, user_prompt):
    """Process image with GPT-4V (placeholder logic)"""
    with open(image_path, "rb") as imgf:
        image_data = imgf.read()
    b64img = base64.b64encode(image_data).decode("utf-8")
    resp = openai_client.chat.completions.create(
        model=st.session_state["openai_model"],
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": user_prompt},
                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64img}"}}
                ]
            }
        ],
        temperature=0.0,
    )
    return resp.choices[0].message.content

def process_audio(audio_path):
    """Process audio with Whisper (placeholder logic)"""
    with open(audio_path, "rb") as f:
        transcription = openai_client.audio.transcriptions.create(model="whisper-1", file=f)
    st.session_state.messages.append({"role": "user", "content": transcription.text})
    return transcription.text

def process_video(video_path, seconds_per_frame=1):
    """Extract frames from video (placeholder logic)"""
    vid = cv2.VideoCapture(video_path)
    total = int(vid.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = vid.get(cv2.CAP_PROP_FPS)
    skip = int(fps * seconds_per_frame)
    frames_b64 = []
    for i in range(0, total, skip):
        vid.set(cv2.CAP_PROP_POS_FRAMES, i)
        ret, frame = vid.read()
        if not ret:
            break
        _, buf = cv2.imencode(".jpg", frame)
        frames_b64.append(base64.b64encode(buf).decode("utf-8"))
    vid.release()
    return frames_b64

def process_video_with_gpt(video_path, prompt):
    """Analyze video frames with GPT-4V (placeholder logic)"""
    frames = process_video(video_path)
    resp = openai_client.chat.completions.create(
        model=st.session_state["openai_model"],
        messages=[
            {"role": "system", "content": "Analyze video frames."},
            {
                "role": "user",
                "content": [
                    {"type":"text","text":prompt},
                    *[{"type":"image_url","image_url":{"url":f"data:image/jpeg;base64,{fr}"}} for fr in frames]
                ]
            }
        ]
    )
    return resp.choices[0].message.content

# 🤖 9. AI Model Integration

def save_full_transcript(query, text):
    """Save full transcript of Arxiv results as a file."""
    create_file(query, text, "md")

# ---------------------------------------------------
# NEW: Extremely simple "parse_arxiv_refs" logic
# that reads each non-empty line, up to 20 lines.
# Extract bracketed title if present, year if present.
# The entire line is the "summary" for display + TTS.
# ---------------------------------------------------
def parse_arxiv_refs(ref_text: str):
    lines = ref_text.split('\n')
    # remove empty lines
    lines = [ln.strip() for ln in lines if ln.strip()]
    # limit to 20
    lines = lines[:20]

    refs = []
    for ln in lines:
        # bracketed title if found
        bracket_match = re.search(r"\[([^\]]+)\]", ln)
        title = bracket_match.group(1) if bracket_match else "No Title"
        # find a year 20xx if present
        year_match = re.search(r"(20\d{2})", ln)
        year = int(year_match.group(1)) if year_match else None
        
        refs.append({
            "line": ln,      # the entire raw line for display
            "title": title,  # bracketed content or "No Title"
            "year": year     # e.g. 2023, 2024, or None
        })
    return refs

def perform_ai_lookup(q, vocal_summary=True, extended_refs=False, 
                      titles_summary=True, full_audio=False):
    """
    1) Query the RAG pipeline
    2) Display results
    3) Also parse references into lines, up to 20
    4) Show each reference with full content
    5) If year in [2023, 2024], auto-generate TTS
    """
    start = time.time()

    # 1) Query HF RAG pipeline
    client = Client("awacke1/Arxiv-Paper-Search-And-QA-RAG-Pattern")
    # 20 references
    refs = client.predict(q, 20, "Semantic Search", "mistralai/Mixtral-8x7B-Instruct-v0.1",
                          api_name="/update_with_rag_md")[0]
    # Main summary
    r2 = client.predict(q, "mistralai/Mixtral-8x7B-Instruct-v0.1", True, api_name="/ask_llm")

    # 2) Combine for final text
    result = f"### 🔎 {q}\n\n{r2}\n\n{refs}"
    st.markdown(result)

    # Optionally produce "all at once" TTS
    if full_audio:
        complete_text = f"Complete response for query: {q}. {clean_for_speech(r2)} {clean_for_speech(refs)}"
        audio_file_full = speak_with_edge_tts(complete_text)
        st.write("### 📚 Full Audio")
        play_and_download_audio(audio_file_full)

    if vocal_summary:
        main_text = clean_for_speech(r2)
        audio_file_main = speak_with_edge_tts(main_text)
        st.write("### 🎙 Short Audio")
        play_and_download_audio(audio_file_main)

    if extended_refs:
        summaries_text = "Extended references: " + refs.replace('"','')
        summaries_text = clean_for_speech(summaries_text)
        audio_file_refs = speak_with_edge_tts(summaries_text)
        st.write("### 📜 Long Refs")
        play_and_download_audio(audio_file_refs)

    # 3) Parse references
    parsed = parse_arxiv_refs(refs)

    # 4) Show references
    st.write("## Individual Paper Lines (Up to 20)")
    for i, ref in enumerate(parsed):
        st.markdown(f"**Ref #{i+1}**:  {ref['line']}")
        if ref['year'] in [2023, 2024]:
            # TTS content: "Title + entire line"
            tts_text = f"Title: {ref['title']}. Full content: {ref['line']}"
            out_fn = generate_audio_filename(q, ref['title'], ref['line'])
            tmp_mp3 = speak_with_edge_tts(tts_text)
            if tmp_mp3 and os.path.exists(tmp_mp3):
                # rename to out_fn
                os.rename(tmp_mp3, out_fn)
                # auto-play
                auto_play_audio(out_fn)
        st.write("---")

    # Titles only block
    if titles_summary:
        # This was your older code - parse bracketed titles from each line
        # to produce an all-in-one TTS if desired
        lines = refs.split('\n')
        titles = []
        for line in lines:
            m = re.search(r"\[([^\]]+)\]", line)
            if m:
                titles.append(m.group(1))
        if titles:
            titles_text = "Titles: " + ", ".join(titles)
            titles_text = clean_for_speech(titles_text)
            audio_file_titles = speak_with_edge_tts(titles_text)
            st.write("### 🔖 Titles (All-In-One)")
            play_and_download_audio(audio_file_titles)

    elapsed = time.time() - start
    st.write(f"**Total Elapsed:** {elapsed:.2f} s")

    # 5) Save entire text as MD file
    create_file(q, result, "md")

    return result

def process_with_gpt(text):
    """Process text with GPT-4"""
    if not text:
        return
    st.session_state.messages.append({"role":"user","content":text})
    with st.chat_message("user"):
        st.markdown(text)
    with st.chat_message("assistant"):
        c = openai_client.chat.completions.create(
            model=st.session_state["openai_model"],
            messages=st.session_state.messages,
            stream=False
        )
        ans = c.choices[0].message.content
        st.write("GPT-4o: " + ans)
        create_file(text, ans, "md")
        st.session_state.messages.append({"role":"assistant","content":ans})
    return ans

def process_with_claude(text):
    """Process text with Claude"""
    if not text:
        return
    with st.chat_message("user"):
        st.markdown(text)
    with st.chat_message("assistant"):
        r = claude_client.messages.create(
            model="claude-3-sonnet-20240229",
            max_tokens=1000,
            messages=[{"role":"user","content":text}]
        )
        ans = r.content[0].text
        st.write("Claude-3.5: " + ans)
        create_file(text, ans, "md")
        st.session_state.chat_history.append({"user":text,"claude":ans})
    return ans

# 📂 10. File Management
def create_zip_of_files(md_files, mp3_files):
    """Create zip with a short naming approach"""
    md_files = [f for f in md_files if os.path.basename(f).lower() != 'readme.md']
    all_files = md_files + mp3_files
    if not all_files:
        return None

    # Collect content for high-info term extraction
    all_content = []
    for f in all_files:
        if f.endswith('.md'):
            with open(f, 'r', encoding='utf-8') as file:
                all_content.append(file.read())
        elif f.endswith('.mp3'):
            all_content.append(os.path.basename(f))
    
    combined_content = " ".join(all_content)
    info_terms = get_high_info_terms(combined_content)
    
    timestamp = datetime.now().strftime("%y%m_%H%M")
    name_text = '_'.join(term.replace(' ', '-') for term in info_terms[:3])
    zip_name = f"{timestamp}_{name_text}.zip"
    
    with zipfile.ZipFile(zip_name,'w') as z:
        for f in all_files:
            z.write(f)
    
    return zip_name

def load_files_for_sidebar():
    """Load and group files for sidebar display"""
    md_files = glob.glob("*.md")
    mp3_files = glob.glob("*.mp3")
    md_files = [f for f in md_files if os.path.basename(f).lower() != 'readme.md']

    all_files = md_files + mp3_files
    groups = defaultdict(list)
    for f in all_files:
        fname = os.path.basename(f)
        prefix = fname[:10]  # e.g. "2310_1205_"
        groups[prefix].append(f)

    for prefix in groups:
        groups[prefix].sort(key=lambda x: os.path.getmtime(x), reverse=True)

    sorted_prefixes = sorted(groups.keys(), 
                             key=lambda pre: max(os.path.getmtime(x) for x in groups[pre]), 
                             reverse=True)
    return groups, sorted_prefixes

def extract_keywords_from_md(files):
    """Extract keywords from markdown files"""
    text = ""
    for f in files:
        if f.endswith(".md"):
            c = open(f,'r',encoding='utf-8').read()
            text += " " + c
    return get_high_info_terms(text)

def display_file_manager_sidebar(groups, sorted_prefixes):
    """Display file manager in sidebar"""
    st.sidebar.title("🎵 Audio & Docs Manager")

    all_md = []
    all_mp3 = []
    for prefix in groups:
        for f in groups[prefix]:
            if f.endswith(".md"):
                all_md.append(f)
            elif f.endswith(".mp3"):
                all_mp3.append(f)

    top_bar = st.sidebar.columns(3)
    with top_bar[0]:
        if st.button("🗑 DelAllMD"):
            for f in all_md:
                os.remove(f)
            st.session_state.should_rerun = True
    with top_bar[1]:
        if st.button("🗑 DelAllMP3"):
            for f in all_mp3:
                os.remove(f)
            st.session_state.should_rerun = True
    with top_bar[2]:
        if st.button("⬇️ ZipAll"):
            z = create_zip_of_files(all_md, all_mp3)
            if z:
                st.sidebar.markdown(get_download_link(z), unsafe_allow_html=True)

    for prefix in sorted_prefixes:
        files = groups[prefix]
        kw = extract_keywords_from_md(files)
        keywords_str = " ".join(kw) if kw else "No Keywords"
        with st.sidebar.expander(f"{prefix} Files ({len(files)}) - KW: {keywords_str}", expanded=True):
            c1,c2 = st.columns(2)
            with c1:
                if st.button("👀ViewGrp", key="view_group_"+prefix):
                    st.session_state.viewing_prefix = prefix
            with c2:
                if st.button("🗑DelGrp", key="del_group_"+prefix):
                    for f in files:
                        os.remove(f)
                    st.success(f"Deleted group {prefix}!")
                    st.session_state.should_rerun = True

            for f in files:
                fname = os.path.basename(f)
                ctime = datetime.fromtimestamp(os.path.getmtime(f)).strftime("%Y-%m-%d %H:%M:%S")
                st.write(f"**{fname}** - {ctime}")

# 🎯 11. Main Application
def main():
    st.sidebar.markdown("### 🚲BikeAI🏆 Multi-Agent Research")
    tab_main = st.radio("Action:", ["🎤 Voice","📸 Media","🔍 ArXiv","📝 Editor"], horizontal=True)

    # If you have a custom React component
    mycomponent = components.declare_component("mycomponent", path="mycomponent")
    val = mycomponent(my_input_value="Hello")

    # Show input in a text box for editing if detected
    if val:
        val_stripped = val.replace('\n', ' ')
        edited_input = st.text_area("✏️ Edit Input:", value=val_stripped, height=100)
        run_option = st.selectbox("Model:", ["Arxiv", "GPT-4o", "Claude-3.5"])
        col1, col2 = st.columns(2)
        with col1:
            autorun = st.checkbox("⚙ AutoRun", value=True)
        with col2:
            full_audio = st.checkbox("📚FullAudio", value=False, 
                                     help="Generate full audio response")

        input_changed = (val != st.session_state.old_val)

        if autorun and input_changed:
            st.session_state.old_val = val
            if run_option == "Arxiv":
                perform_ai_lookup(edited_input,
                                  vocal_summary=True,
                                  extended_refs=False, 
                                  titles_summary=True,
                                  full_audio=full_audio)
            elif run_option == "GPT-4o":
                process_with_gpt(edited_input)
            elif run_option == "Claude-3.5":
                process_with_claude(edited_input)
        else:
            if st.button("▶ Run"):
                st.session_state.old_val = val
                if run_option == "Arxiv":
                    perform_ai_lookup(edited_input,
                                      vocal_summary=True,
                                      extended_refs=False, 
                                      titles_summary=True,
                                      full_audio=full_audio)
                elif run_option == "GPT-4o":
                    process_with_gpt(edited_input)
                elif run_option == "Claude-3.5":
                    process_with_claude(edited_input)

    if tab_main == "🔍 ArXiv":
        st.subheader("🔍 Query ArXiv")
        q = st.text_input("🔍 Query:")

        st.markdown("### 🎛 Options")
        vocal_summary = st.checkbox("🎙ShortAudio", value=True)
        extended_refs = st.checkbox("📜LongRefs", value=False)
        titles_summary = st.checkbox("🔖TitlesOnly", value=True)
        full_audio = st.checkbox("📚FullAudio", value=False,
                                 help="Full audio of results")
        full_transcript = st.checkbox("🧾FullTranscript", value=False,
                                      help="Generate a full transcript file")

        if q and st.button("🔍Run"):
            result = perform_ai_lookup(q, 
                                       vocal_summary=vocal_summary,
                                       extended_refs=extended_refs, 
                                       titles_summary=titles_summary,
                                       full_audio=full_audio)
            if full_transcript:
                save_full_transcript(q, result)

        st.markdown("### Change Prompt & Re-Run")
        q_new = st.text_input("🔄 Modify Query:")
        if q_new and st.button("🔄 Re-Run with Modified Query"):
            result = perform_ai_lookup(q_new, 
                                       vocal_summary=vocal_summary,
                                       extended_refs=extended_refs, 
                                       titles_summary=titles_summary,
                                       full_audio=full_audio)
            if full_transcript:
                save_full_transcript(q_new, result)

    elif tab_main == "🎤 Voice":
        st.subheader("🎤 Voice Input")
        user_text = st.text_area("💬 Message:", height=100)
        user_text = user_text.strip().replace('\n', ' ')
        if st.button("📨 Send"):
            process_with_gpt(user_text)
        st.subheader("📜 Chat History")
        t1, t2 = st.tabs(["Claude History","GPT-4o History"])
        with t1:
            for c in st.session_state.chat_history:
                st.write("**You:**", c["user"])
                st.write("**Claude:**", c["claude"])
        with t2:
            for m in st.session_state.messages:
                with st.chat_message(m["role"]):
                    st.markdown(m["content"])

    elif tab_main == "📸 Media":
        st.header("📸 Images & 🎥 Videos")
        tabs = st.tabs(["🖼 Images", "🎥 Video"])
        with tabs[0]:
            imgs = glob.glob("*.png")+glob.glob("*.jpg")
            if imgs:
                c = st.slider("Cols", 1, 5, 3)
                cols = st.columns(c)
                for i, f in enumerate(imgs):
                    with cols[i % c]:
                        st.image(Image.open(f), use_container_width=True)
                        if st.button(f"👀 Analyze {os.path.basename(f)}", key=f"analyze_{f}"):
                            a = process_image(f, "Describe this image.")
                            st.markdown(a)
            else:
                st.write("No images found.")
        with tabs[1]:
            vids = glob.glob("*.mp4")
            if vids:
                for v in vids:
                    with st.expander(f"🎥 {os.path.basename(v)}"):
                        st.video(v)
                        if st.button(f"Analyze {os.path.basename(v)}", key=f"analyze_{v}"):
                            a = process_video_with_gpt(v, "Describe video.")
                            st.markdown(a)
            else:
                st.write("No videos found.")

    elif tab_main == "📝 Editor":
        if getattr(st.session_state, 'current_file', None):
            st.subheader(f"Editing: {st.session_state.current_file}")
            new_text = st.text_area("✏️ Content:", st.session_state.file_content, height=300)
            if st.button("💾 Save"):
                with open(st.session_state.current_file, 'w', encoding='utf-8') as f:
                    f.write(new_text)
                st.success("Updated!")
                st.session_state.should_rerun = True
        else:
            st.write("Select a file from the sidebar to edit.")

    # File manager in sidebar
    groups, sorted_prefixes = load_files_for_sidebar()
    display_file_manager_sidebar(groups, sorted_prefixes)

    # If user clicked "view group"
    if st.session_state.viewing_prefix and st.session_state.viewing_prefix in groups:
        st.write("---")
        st.write(f"**Viewing Group:** {st.session_state.viewing_prefix}")
        for f in groups[st.session_state.viewing_prefix]:
            fname = os.path.basename(f)
            ext = os.path.splitext(fname)[1].lower().strip('.')
            st.write(f"### {fname}")
            if ext == "md":
                content = open(f, 'r', encoding='utf-8').read()
                st.markdown(content)
            elif ext == "mp3":
                st.audio(f)
            else:
                st.markdown(get_download_link(f), unsafe_allow_html=True)
        if st.button("❌ Close"):
            st.session_state.viewing_prefix = None

    if st.session_state.should_rerun:
        st.session_state.should_rerun = False
        st.rerun()

if __name__ == "__main__":
    main()