import streamlit as st
import anthropic
import openai
import base64
import cv2
import glob
import os
import re
import asyncio
import edge_tts
from datetime import datetime
from collections import defaultdict
from dotenv import load_dotenv
from gradio_client import Client
from PIL import Image

# 🎯 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', "")
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)

# 📝 3. Session State Management
if 'parsed_papers' not in st.session_state:
    st.session_state['parsed_papers'] = []
if 'audio_generated' not in st.session_state:
    st.session_state['audio_generated'] = {}
if 'voices' not in st.session_state:
    st.session_state['voices'] = []
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

# 🎨 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(prefix, title, file_type="md"):
    """
    Generate filename with meaningful terms and prefix.
    The filename includes a timestamp and a cleaned title.
    """
    timestamp = datetime.now().strftime("%y%m_%H%M")
    title_cleaned = clean_text_for_filename(title)
    filename = f"{timestamp}_{prefix}_{title_cleaned}.{file_type}"
    return filename

def create_md_file(paper):
    """Create Markdown file for a paper."""
    filename = generate_filename("paper", paper['title'], "md")
    content = f"# {paper['title']}\n\n**Year:** {paper['year'] if paper['year'] else 'Unknown'}\n\n**Summary:**\n{paper['summary']}"
    with open(filename, 'w', encoding='utf-8') as f:
        f.write(content)
    return filename

def get_download_link(file):
    """Generate download link for file."""
    with open(file, "rb") as f_file:
        b64 = base64.b64encode(f_file.read()).decode()
    mime_type = "audio/mpeg" if file.endswith(".mp3") else "text/markdown"
    return f'<a href="data:{mime_type};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

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("audio", text[:50], "mp3")
    await communicate.save(out_fn)
    return out_fn

def speak_with_edge_tts(text, voice, rate=0, pitch=0):
    """Wrapper for Edge TTS generation."""
    try:
        return asyncio.run(edge_tts_generate_audio(text, voice, rate, pitch))
    except Exception as e:
        st.error(f"Error generating audio: {e}")
        return None

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 = get_download_link(file_path)
        st.markdown(dl_link, unsafe_allow_html=True)

# 🎬 8. 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.ChatCompletion.create(
        model=st.session_state["openai_model"],
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": f"{user_prompt} Image data: data:image/png;base64,{b64img}"}
        ],
        temperature=0.0,
    )
    return resp.choices[0].message.content

def process_audio_file(audio_path):
    """Process audio with Whisper."""
    with open(audio_path, "rb") as f:
        transcription = openai.Audio.transcribe("whisper-1", f)
    return transcription['text']

def process_video(video_path, seconds_per_frame=1):
    """Extract frames from video."""
    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)
    combined_images = " ".join([f"data:image/jpeg;base64,{fr}" for fr in frames])
    resp = openai.ChatCompletion.create(
        model=st.session_state["openai_model"],
        messages=[
            {"role":"system","content":"Analyze the following video frames."},
            {"role":"user","content": f"{prompt} Frames: {combined_images}"}
        ]
    )
    return resp.choices[0].message.content

# 🤖 9. AI Model Integration
def parse_papers(transcript_text: str):
    """
    Parse the transcript text into individual papers.
    Assumes that each paper starts with a number and is enclosed in brackets for the title and year.
    Example:
        1) [Paper Title (2023)] This is the summary...
    """
    papers = []
    # Split based on numbered entries
    paper_blocks = re.split(r'\d+\)\s*\[', transcript_text)
    for block in paper_blocks[1:]:  # Skip the first split as it doesn't contain paper info
        try:
            title_year, summary = block.split(']', 1)
            # Extract title and year using regex
            title_match = re.match(r"(.+?)\s*\((\d{4})\)", title_year)
            if title_match:
                title = title_match.group(1).strip()
                year = int(title_match.group(2))
            else:
                title = title_year.strip()
                year = None
            summary = summary.strip()
            papers.append({
                'title': title,
                'year': year,
                'summary': summary
            })
        except ValueError:
            continue  # Skip blocks that don't match the expected format
    return papers

def save_paper_files(paper, voice):
    """Generate and save Markdown and MP3 files for a paper."""
    # Create Markdown file
    md_filename = create_md_file(paper)
    
    # Generate audio for the entire paper
    audio_text = f"{paper['title']}. {paper['summary']}"
    audio_filename = speak_with_edge_tts(audio_text, voice)
    
    return md_filename, audio_filename

def display_papers(papers, voice):
    """Display all papers with options to generate audio."""
    for idx, paper in enumerate(papers):
        st.markdown(f"### {idx + 1}. {paper['title']} ({paper['year'] if paper['year'] else 'Unknown Year'})")
        st.markdown(f"**Summary:** {paper['summary']}")
        
        # Button to generate and play audio
        if st.button(f"🔊 Read Aloud - {paper['title']}", key=f"read_aloud_{idx}"):
            md_file, audio_file = save_paper_files(paper, voice)
            if audio_file:
                st.success("Audio generated successfully!")
                play_and_download_audio(audio_file)
            else:
                st.error("Failed to generate audio.")
        
        st.write("---")

def cache_parsed_papers(papers):
    """Cache the parsed papers."""
    st.session_state['parsed_papers'] = papers

def get_cached_papers():
    """Retrieve cached papers."""
    return st.session_state.get('parsed_papers', [])

def save_full_transcript(query, text):
    """Save full transcript of Arxiv results as a file."""
    filename = generate_filename("transcript", query, "md")
    with open(filename, 'w', encoding='utf-8') as f:
        f.write(text)
    return filename

def perform_ai_lookup(q, vocal_summary=True, extended_refs=False, 
                      titles_summary=True, full_audio=False, selected_voice="en-US-AriaNeural"):
    """Perform Arxiv search and generate audio summaries."""
    start = time.time()

    # 🎯 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
    clean_q = q.replace('\n', ' ')
    result = f"### 🔎 {clean_q}\n\n{r2}\n\n{refs}"
    st.markdown(result)
    
    # 🎯 3) Parse papers from the references
    parsed_papers = parse_papers(refs)
    cache_parsed_papers(parsed_papers)
    
    # 🎯 4) Display all parsed papers with options
    st.write("## Individual Papers")
    display_papers(parsed_papers, selected_voice)
    
    elapsed = time.time() - start
    st.write(f"**Total Elapsed:** {elapsed:.2f} s")
    
    # Always create a file with the result
    save_full_transcript(clean_q, result)
    
    return result

# 📂 10. 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]
        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("👀 View Group", key="view_group_"+prefix):
                    st.session_state.viewing_prefix = prefix
            with c2:
                if st.button("🗑 Delete Group", 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
async def get_available_voices():
    voices = await edge_tts.list_voices()
    return [voice["ShortName"] for voice in voices if voice["Locale"].startswith("en")]

@st.cache_resource
def fetch_voices():
    return asyncio.run(get_available_voices())

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

    # Initialize voices if not already done
    if not st.session_state['voices']:
        st.session_state['voices'] = fetch_voices()

    st.sidebar.markdown("### 🎤 Select Voice for Audio Generation")
    selected_voice = st.sidebar.selectbox(
        "Choose a voice:",
        options=st.session_state['voices'],
        index=st.session_state['voices'].index("en-US-AriaNeural") if "en-US-AriaNeural" in st.session_state['voices'] else 0
    )

    # Main Tabs
    if tab_main == "🔍 ArXiv":
        st.subheader("🔍 Query ArXiv")
        q = st.text_input("🔍 Query:").replace('\n', ' ')

        st.markdown("### 🎛 Options")
        vocal_summary = st.checkbox("🎙 Short Audio", value=True)
        extended_refs = st.checkbox("📜 Long References", value=False)
        titles_summary = st.checkbox("🔖 Titles Only", value=True)
        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"):
            result = perform_ai_lookup(
                q, 
                vocal_summary=vocal_summary, 
                extended_refs=extended_refs, 
                titles_summary=titles_summary, 
                full_audio=full_audio, 
                selected_voice=selected_voice
            )
            if full_transcript:
                save_full_transcript(q, result)

        st.markdown("### Change Prompt & Re-Run")
        q_new = st.text_input("🔄 Modify Query:").replace('\n', ' ')
        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, 
                selected_voice=selected_voice
            )
            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.get('chat_history', []):
                st.write("**You:**", c["user"])
                st.write("**Claude:**", c["claude"])
        with t2:
            for m in st.session_state.get('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:
                cols = st.columns(min(5, len(imgs)))
                for i, f in enumerate(imgs[:20]):
                    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}"):
                            analysis = process_image(f, "Describe this image.")
                            st.markdown(analysis)
            else:
                st.write("No images found.")

        with tabs[1]:
            vids = glob.glob("*.mp4")[:20]
            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}"):
                            analysis = process_video_with_gpt(v, "Describe video.")
                            st.markdown(analysis)
            else:
                st.write("No videos found.")

    elif tab_main == "📝 Editor":
        st.subheader("📝 Editor")
        files = glob.glob("*.md")
        if files:
            selected_file = st.selectbox("Select a file to edit:", files)
            if selected_file:
                with open(selected_file, 'r', encoding='utf-8') as f:
                    file_content = f.read()
                new_text = st.text_area("✏️ Content:", file_content, height=300)
                if st.button("💾 Save"):
                    with open(selected_file, 'w', encoding='utf-8') as f:
                        f.write(new_text)
                    st.success("File updated successfully!")
                    st.session_state.should_rerun = True
        else:
            st.write("No Markdown files available for editing.")

    # File Manager Sidebar
    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(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 process_with_gpt(text):
    """Process text with GPT-4."""
    if not text: 
        return
    # Initialize messages if not present
    if 'messages' not in st.session_state:
        st.session_state['messages'] = []
    st.session_state['messages'].append({"role":"user","content":text})
    with st.chat_message("user"):
        st.markdown(text)
    with st.chat_message("assistant"):
        try:
            response = openai.ChatCompletion.create(
                model=st.session_state["openai_model"],
                messages=st.session_state['messages'],
                stream=False
            )
            ans = response.choices[0].message.content
            st.write("GPT-4o: " + ans)
            create_md_file({"title": "User Query", "year": None, "summary": ans})
            st.session_state['messages'].append({"role":"assistant","content":ans})
        except Exception as e:
            st.error(f"Error processing with GPT-4: {e}")

def process_with_claude(text):
    """Process text with Claude."""
    if not text: 
        return
    # Initialize chat_history if not present
    if 'chat_history' not in st.session_state:
        st.session_state['chat_history'] = []
    with st.chat_message("user"):
        st.markdown(text)
    with st.chat_message("assistant"):
        try:
            response = claude_client.messages.create(
                model="claude-3-sonnet-20240229",
                max_tokens=1000,
                messages=[{"role":"user","content":text}]
            )
            ans = response.content[0].text
            st.write("Claude-3.5: " + ans)
            create_md_file({"title": "User Query", "year": None, "summary": ans})
            st.session_state['chat_history'].append({"user":text,"claude":ans})
        except Exception as e:
            st.error(f"Error processing with Claude: {e}")

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