import streamlit as st
import anthropic
import openai
import base64
import os
import re
import asyncio
from datetime import datetime
from gradio_client import Client
from collections import defaultdict
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"
    }
)
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)

# 🔑 2. API Setup & Clients
from dotenv import load_dotenv
load_dotenv()

openai_api_key = os.getenv('OPENAI_API_KEY', "")
anthropic_key = os.getenv('ANTHROPIC_API_KEY_3', "")
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  # Using OpenAI directly

# 📝 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-4"  # Update as needed
if 'messages' not in st.session_state:
    st.session_state['messages'] = []
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. 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]

# 📁 5. 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>'

# 🔊 6. 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

async def edge_tts_generate_audio(text, voice="en-US-AriaNeural", rate=0, pitch=0, out_fn="temp.mp3"):
    """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)
    await communicate.save(out_fn)
    return out_fn

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

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):
    """Embeds an <audio> tag with autoplay + controls + a download link."""
    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):
    """
    Generate a specialized MP3 filename using query + title + summary.
    Example: "2310_1205_query_title_summary.mp3"
    """
    combined = (query + " " + title + " " + summary).strip().lower()
    combined = re.sub(r'[^\w\s-]', '', combined)  # Remove special characters
    combined = "_".join(combined.split())[:80]    # Limit length
    prefix = datetime.now().strftime("%y%m_%H%M")
    return f"{prefix}_{combined}.mp3"

# 🎬 7. 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")
    resp = openai_client.ChatCompletion.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_with_whisper(audio_path):
    """Process audio with Whisper"""
    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"""
    import cv2
    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"""
    frames = process_video(video_path)
    resp = openai_client.ChatCompletion.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

# 🤖 8. AI Model Integration
def save_full_transcript(query, text):
    """Save full transcript of Arxiv results as a file."""
    create_file(query, text, "md")

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.ChatCompletion.create(
            model=st.session_state["openai_model"],
            messages=st.session_state.messages,
            stream=False
        )
        ans = c.choices[0].message.content
        st.write("GPT-4: " + 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.completions.create(
            prompt=text,
            model="claude-3",
            max_tokens=1000
        )
        ans = r['completion']
        st.write("Claude-3.5: " + ans)
        create_file(text, ans, "md")
        st.session_state.chat_history.append({"user":text,"claude":ans})
    return ans

# 📂 9. 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

    # 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}")

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

    # Placeholder for custom component if needed
    # mycomponent = components.declare_component("mycomponent", path="mycomponent")
    # val = mycomponent(my_input_value="Hello")

    # Example input handling
    # if val:
    #     # Handle custom component input
    #     pass

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

        st.markdown("### 🎛 Options")
        full_audio = st.checkbox("📚 Full Audio", value=False, help="Generate full audio response")
        full_transcript = st.checkbox("🧾 Full Transcript", value=False, help="Generate a full transcript file")

        if q and st.button("🔍 Run Query"):
            perform_ai_lookup(q)
            if full_transcript:
                create_file(q, "Full transcript generated.", "md")  # Customize as needed

    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-4 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") + glob.glob("*.jpeg")
            if imgs:
                cols = st.columns(st.slider("Cols", 1, 5, 3))
                for i, f in enumerate(imgs):
                    with cols[i % len(cols)]:
                        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") + glob.glob("*.avi") + glob.glob("*.mov")
            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 this video.")
                            st.markdown(a)
            else:
                st.write("No videos found.")

    elif tab_main == "📝 Editor":
        st.subheader("📝 File Editor")
        # Example editor logic: list markdown files and allow editing
        md_files = glob.glob("*.md")
        if md_files:
            selected_file = st.selectbox("Select a file to edit:", md_files)
            with st.form("edit_form"):
                new_content = st.text_area("✏️ Content:", open(selected_file, 'r', encoding='utf-8').read(), height=300)
                submitted = st.form_submit_button("💾 Save")
                if submitted:
                    with open(selected_file, 'w', encoding='utf-8') as f:
                        f.write(new_content)
                    st.success(f"Updated {selected_file}!")
                    st.session_state.should_rerun = True
        else:
            st.write("No markdown files available 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.experimental_rerun()

def parse_arxiv_papers(ref_text: str):
    """
    Splits the references into paper-level chunks.
    Each paper starts with a number followed by a parenthesis, e.g., "1) [Title (Year)] Summary..."
    Returns a list of dictionaries with 'title', 'summary', and 'year'.
    Limits to 20 papers.
    """
    # Split based on patterns like "1) ", "2) ", etc.
    chunks = re.split(r'\n?\d+\)\s+', ref_text)
    # Remove any empty strings resulting from split
    chunks = [chunk.strip() for chunk in chunks if chunk.strip()]
    papers = []
    for chunk in chunks[:20]:
        # Extract title within brackets if present
        title_match = re.search(r'\[([^\]]+)\]', chunk)
        title = title_match.group(1).strip() if title_match else "No Title"
        
        # Extract year (assuming it's a 4-digit number within the title or summary)
        year_match = re.search(r'\b(20\d{2})\b', chunk)
        year = int(year_match.group(1)) if year_match else None
        
        # The entire chunk is considered the summary
        summary = chunk
        
        papers.append({
            'title': title,
            'summary': summary,
            'year': year
        })
    return papers

def perform_ai_lookup(q):
    """
    Performs the Arxiv search and handles the processing of results.
    Generates audio files for each paper (if year is 2023 or 2024).
    """
    st.write(f"## Query: {q}")

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

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

    # 3) Parse references into papers
    papers = parse_arxiv_papers(refs)

    # 4) Display each paper and generate audio if applicable
    st.write("## Individual Papers (Up to 20)")
    for idx, paper in enumerate(papers):
        year_str = paper["year"] if paper["year"] else "Unknown Year"
        st.markdown(f"**Paper #{idx+1}: {paper['title']}**  \n*Year:* {year_str}")
        st.markdown(f"*Summary:* {paper['summary']}")
        st.write("---")

        # Generate TTS if year is 2023 or 2024
        if paper["year"] in [2023, 2024]:
            # Combine title and summary for TTS
            tts_text = f"Title: {paper['title']}. Summary: {paper['summary']}"
            # Generate a specialized filename
            mp3_filename = generate_audio_filename(q, paper['title'], paper['summary'])
            # Generate audio using Edge TTS
            temp_mp3 = speak_with_edge_tts(tts_text, out_fn=mp3_filename)
            if temp_mp3 and os.path.exists(mp3_filename):
                # Embed the audio player with auto-play and download link
                auto_play_audio(mp3_filename)

    # Optionally save the full transcript
    st.write("### Transcript")
    st.markdown(result)
    create_file(q, result, "md")

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.ChatCompletion.create(
            model=st.session_state["openai_model"],
            messages=st.session_state.messages,
            stream=False
        )
        ans = c.choices[0].message.content
        st.write("GPT-4: " + 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.completions.create(
            prompt=text,
            model="claude-3",
            max_tokens=1000
        )
        ans = r['completion']
        st.write("Claude-3.5: " + ans)
        create_file(text, ans, "md")
        st.session_state.chat_history.append({"user":text,"claude":ans})
    return ans

# Run the app
if __name__ == "__main__":
    main()