import streamlit as st
import anthropic, openai, base64, cv2, glob, json, math, os, pytz, random, re, requests, textract, time, 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. Text Cleaning Functionality
class TextCleaner:
    """Helper class for text cleaning operations"""
    def __init__(self):
        self.replacements = {
            "\\n": " ",     # Replace escaped newlines
            "</s>": "",     # Remove end tags
            "<s>": "",      # Remove start tags
            "\n": " ",      # Replace actual newlines
            "\r": " ",      # Replace carriage returns
            "\t": " ",      # Replace tabs
        }
        
        self.preserve_replacements = {
            "\\n": "\n",    # Convert escaped to actual newlines
            "</s>": "",     # Remove end tags
            "<s>": "",      # Remove start tags
            "\r": "\n",     # Convert returns to newlines
            "\t": "    "    # Convert tabs to spaces
        }

    def clean_text(self, text: str, preserve_format: bool = False) -> str:
        """
        Clean text removing problematic characters and normalizing whitespace.
        Args:
            text: Text to clean
            preserve_format: Whether to preserve some formatting (newlines etc)
        Returns:
            Cleaned text string
        """
        if not text or not isinstance(text, str):
            return ""
            
        replacements = (self.preserve_replacements if preserve_format 
                       else self.replacements)
        
        cleaned = text
        for old, new in replacements.items():
            cleaned = cleaned.replace(old, new)
            
        # Normalize whitespace while preserving paragraphs if needed
        if preserve_format:
            cleaned = re.sub(r'\n{3,}', '\n\n', cleaned)
        else:
            cleaned = re.sub(r'\s+', ' ', cleaned)
            
        return cleaned.strip()

    def clean_dict(self, data: dict, fields: list) -> dict:
        """Clean specified fields in a dictionary"""
        if not data or not isinstance(data, dict):
            return {}
            
        cleaned = data.copy()
        for field in fields:
            if field in cleaned:
                cleaned[field] = self.clean_text(cleaned[field])
        return cleaned

    def clean_list(self, items: list, fields: list) -> list:
        """Clean specified fields in a list of dictionaries"""
        if not isinstance(items, list):
            return []
        return [self.clean_dict(item, fields) for item in items]

# Initialize cleaner
cleaner = TextCleaner()

# 🔑 3. 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')

# 📝 4. 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-4-1106-preview"
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

# 🎨 5. 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": "🎵",
}

# 🧠 6. High-Information Content Extraction
def get_high_info_terms(text: str) -> list:
    """Extract high-information terms from text, including key phrases."""
    text = cleaner.clean_text(text)
    
    # ... rest of function remains the same ...
    [Your existing get_high_info_terms implementation]

def clean_text_for_filename(text: str) -> str:
    """Remove punctuation and short filler words, return a compact string."""
    text = cleaner.clean_text(text)
    
    # ... rest of function remains the same ...
    [Your existing clean_text_for_filename implementation]

# 📁 7. File Operations
def generate_filename(prompt, response, file_type="md"):
    """Generate filename with meaningful terms."""
    cleaned_prompt = cleaner.clean_text(prompt)
    cleaned_response = cleaner.clean_text(response)
    
    prefix = datetime.now().strftime("%y%m_%H%M") + "_"
    combined = (cleaned_prompt + " " + cleaned_response).strip()
    info_terms = get_high_info_terms(combined)
    
    snippet = (cleaned_prompt[:100] + " " + cleaned_response[:100]).strip()
    snippet_cleaned = clean_text_for_filename(snippet)
    
    name_parts = info_terms + [snippet_cleaned]
    full_name = '_'.join(name_parts)

    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 intelligent naming"""
    filename = generate_filename(prompt.strip(), response.strip(), file_type)
    
    cleaned_prompt = cleaner.clean_text(prompt)
    cleaned_response = cleaner.clean_text(response, preserve_format=True)
    
    with open(filename, 'w', encoding='utf-8') as f:
        f.write(cleaned_prompt + "\n\n" + cleaned_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>'

# 🔊 8. Audio Processing
def clean_for_speech(text: str) -> str:
    """Clean text for speech synthesis"""
    text = cleaner.clean_text(text)
    text = re.sub(r"\(https?:\/\/[^\)]+\)", "", text)
    return text

@st.cache_resource
def speech_synthesis_html(result):
    """Create HTML for speech synthesis"""
    cleaned_result = clean_for_speech(result)
    html_code = f"""
    <html><body>
    <script>
    var msg = new SpeechSynthesisUtterance("{cleaned_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"""
    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"""
    return asyncio.run(edge_tts_generate_audio(text, voice, rate, pitch))

def play_and_download_audio(file_path):
    """Play and provide 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)

# 🎬 9. Media Processing
def process_image(image_path, user_prompt):
    """Process image with GPT-4V"""
    with open(image_path, "rb") as imgf:
        image_data = imgf.read()
    b64img = base64.b64encode(image_data).decode("utf-8")
    
    cleaned_prompt = cleaner.clean_text(user_prompt)
    
    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": cleaned_prompt},
                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64img}"}}
            ]}
        ],
        temperature=0.0,
    )
    return cleaner.clean_text(resp.choices[0].message.content, preserve_format=True)

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

def process_video(video_path, seconds_per_frame=1):
    """Extract frames from video"""
    # ... function remains the same as it handles binary data ...
    [Your existing process_video implementation]

def process_video_with_gpt(video_path, prompt):
    """Analyze video frames with GPT-4V"""
    frames = process_video(video_path)
    cleaned_prompt = cleaner.clean_text(prompt)
    
    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":cleaned_prompt},
                *[{"type":"image_url","image_url":{"url":f"data:image/jpeg;base64,{fr}"}} 
                  for fr in frames]
            ]}
        ]
    )
    return cleaner.clean_text(resp.choices[0].message.content, preserve_format=True)

# 🤖 10. AI Model Integration
def process_with_claude(text):
    """Process text with Claude"""
    if not text: return
    
    cleaned_input = cleaner.clean_text(text)
    with st.chat_message("user"):
        st.markdown(cleaned_input)
        
    with st.chat_message("assistant"):
        r = claude_client.messages.create(
            model="claude-3-sonnet-20240229",
            max_tokens=1000,
            messages=[{"role":"user","content":cleaned_input}]
        )
        raw_response = r.content[0].text
        cleaned_response = cleaner.clean_text(raw_response, preserve_format=True)
        
        st.write("Claude-3.5: " + cleaned_response)
        create_file(cleaned_input, cleaned_response, "md")
        st.session_state.chat_history.append({
            "user": cleaned_input,
            "claude": cleaned_response
        })
    return cleaned_response

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

def perform_ai_lookup(q, vocal_summary=True, extended_refs=False, titles_summary=True, full_audio=False):
    """Perform Arxiv search and generate audio summaries"""
    cleaned_query = cleaner.clean_text(q)
    start = time.time()
    
    client = Client("awacke1/Arxiv-Paper-Search-And-QA-RAG-Pattern")
    refs = client.predict(cleaned_query, 20, "Semantic Search",
                         "mistralai/Mixtral-8x7B-Instruct-v0.1",
                         api_name="/update_with_rag_md")[0]
    r2 = client.predict(cleaned_query, "mistralai/Mixtral-8x7B-Instruct-v0.1",
                       True, api_name="/ask_llm")

    # Clean responses
    cleaned_r2 = cleaner.clean_text(r2, preserve_format=True)
    cleaned_refs = cleaner.clean_text(refs, preserve_format=True)
    
    result = f"### 🔎 {cleaned_query}\n\n{cleaned_r2}\n\n{cleaned_refs}"
    st.markdown(result)
    
    if full_audio:
        complete_text = f"Complete response for query: {cleaned_query}. {clean_for_speech(cleaned_r2)} {clean_for_speech(cleaned_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(cleaned_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: " + cleaned_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)

    if titles_summary:
        titles = []
        for line in cleaned_refs.split('\n'):
            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")
            play_and_download_audio(audio_file_titles)

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

    create_file(cleaned_query, result, "md")
    return result

def save_full_transcript(query, text):
    """Save full transcript of results as a file."""
    cleaned_query = cleaner.clean_text(query)
    cleaned_text = cleaner.clean_text(text, preserve_format=True)
    create_file(cleaned_query, cleaned_text, "md")

# 📂 11. File Management
def create_zip_of_files(md_files, mp3_files):
    """Create zip with intelligent naming"""
    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

    all_content = []
    for f in all_files:
        if f.endswith('.md'):
            with open(f, 'r', encoding='utf-8') as file:
                content = file.read()
                cleaned_content = cleaner.clean_text(content)
                all_content.append(cleaned_content)
        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]
        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"):
            with open(f, 'r', encoding='utf-8') as file:
                content = file.read()
                cleaned_content = cleaner.clean_text(content)
                text += " " + cleaned_content
    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}")

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

    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:
        cleaned_val = cleaner.clean_text(val)
        edited_input = st.text_area("✏️ Edit Input:", value=cleaned_val, 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)
            else:
                if 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)
                else:
                    if 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:")
        q = cleaner.clean_text(q)

        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="Generate full audio response")
        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:")
        q_new = cleaner.clean_text(q_new)
        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 = cleaner.clean_text(user_text)
        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:**", cleaner.clean_text(c["user"]))
                st.write("**Claude:**", cleaner.clean_text(c["claude"], preserve_format=True))
        with t2:
            for m in st.session_state.messages:
                with st.chat_message(m["role"]):
                    if m["role"] == "user":
                        st.markdown(cleaner.clean_text(m["content"]))
                    else:
                        st.markdown(cleaner.clean_text(m["content"], preserve_format=True))

    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(cleaner.clean_text(a, preserve_format=True))
            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(cleaner.clean_text(a, preserve_format=True))
            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}")
            with open(st.session_state.current_file, 'r', encoding='utf-8') as f:
                content = f.read()
            content = cleaner.clean_text(content, preserve_format=True)
            new_text = st.text_area("✏️ Content:", content, height=300)
            if st.button("💾 Save"):
                cleaned_content = cleaner.clean_text(new_text, preserve_format=True)
                with open(st.session_state.current_file, 'w', encoding='utf-8') as f:
                    f.write(cleaned_content)
                st.success("Updated!")
                st.session_state.should_rerun = True
        else:
            st.write("Select a file from the sidebar to edit.")

    groups, sorted_prefixes = load_files_for_sidebar()
    display_file_manager_sidebar(groups, sorted_prefixes)

    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":
                with open(f, 'r', encoding='utf-8') as file:
                    content = file.read()
                st.markdown(cleaner.clean_text(content, preserve_format=True))
            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()