import streamlit as st
import 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
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
import extra_streamlit_components as stx
from streamlit.runtime.scriptrunner import get_script_run_ctx
import asyncio
import edge_tts

# -------------------- Configuration --------------------
st.set_page_config(
    page_title="🚲CCCG🏆 Code Competition Claude vs GPT",
    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': "🚲CCCG🏆 Code Competition Claude vs GPT"
    }
)
load_dotenv()

# -------------------- Constants --------------------
USER_NAMES = [
    "Aria", "Guy", "Sonia", "Tony", "Jenny", "Davis", "Libby", "Clara", "Liam", "Natasha", "William"
]

ENGLISH_VOICES = [
    "en-US-AriaNeural", "en-US-GuyNeural", "en-GB-SoniaNeural", "en-GB-TonyNeural",
    "en-US-JennyNeural", "en-US-DavisNeural", "en-GB-LibbyNeural", "en-CA-ClaraNeural",
    "en-CA-LiamNeural", "en-AU-NatashaNeural", "en-AU-WilliamNeural"
]

USER_VOICES = dict(zip(USER_NAMES, ENGLISH_VOICES))

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

# -------------------- Session State Initialization --------------------
if 'user_name' not in st.session_state:
    st.session_state['user_name'] = USER_NAMES[0]
if 'old_val' not in st.session_state:
    st.session_state['old_val'] = None
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 'use_streaming' not in st.session_state:
    st.session_state['use_streaming'] = True

# -------------------- Helper Functions --------------------
def get_high_info_terms(text: str) -> list:
    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', 'as', 'if', 'while'
    ])

    key_phrases = [
        'artificial intelligence', 'machine learning', 'deep learning', 'neural networks',
        'natural language processing', 'healthcare systems', 'clinical medicine',
        'genomics', 'biological systems', 'cognitive science', 'data visualization',
        'wellness technology', 'robotics', 'medical imaging', 'semantic understanding',
        'transformers', 'large language models', 'empirical studies', 'scientific research',
        'quantum mechanics', 'biomedical engineering', 'computational biology'
    ]

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

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

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

    return unique_terms[:5]

def clean_text_for_filename(text: str) -> str:
    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]

def generate_filename(prompt, response, file_type="md"):
    central_tz = pytz.timezone('America/Chicago')
    central_time = datetime.now(central_tz)
    prefix = central_time.strftime("%m-%d-%y_%I-%M-%p_")

    combined = (prompt + " " + response).strip()
    info_terms = get_high_info_terms(combined)

    snippet = (prompt[:100] + " " + 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"):
    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):
    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>'

def clean_for_speech(text: str) -> str:
    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

# -------------------- Audio Functions --------------------
async def edge_tts_generate_audio(text, voice="en-US-AriaNeural", rate=0, pitch=0):
    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")
    try:
        await communicate.save(out_fn)
    except edge_tts.exceptions.NoAudioReceived:
        st.error("No audio was received from TTS service.")
        return None
    return out_fn

def speak_with_edge_tts(text, voice="en-US-AriaNeural", rate=0, pitch=0):
    return asyncio.run(edge_tts_generate_audio(text, voice, rate, pitch))

def play_and_download_audio(file_path):
    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)

# -------------------- File Management Functions --------------------
def load_files_for_sidebar():
    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[:17]
        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):
    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 create_zip_of_files(md_files, mp3_files):
    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:
                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 display_file_manager_sidebar(groups, sorted_prefixes):
    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}")

# -------------------- xAI API Functions --------------------
def call_xai_api_batch(query: str) -> dict:
    """
    Call the xAI API in batch mode for complete responses.
    """
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {os.environ.get('xai')}"
    }
    
    data = {
        "messages": [
            {
                "role": "system",
                "content": "You are a helpful scientific research assistant. Analyze the following research query and provide initial insights."
            },
            {
                "role": "user",
                "content": query
            }
        ],
        "model": "grok-2-1212",
        "stream": False,
        "temperature": 0.7
    }
    
    try:
        response = requests.post(
            "https://api.x.ai/v1/chat/completions",
            headers=headers,
            json=data,
            timeout=30
        )
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        st.error(f"Error in batch xAI API call: {str(e)}")
        return None

def stream_xai_response(query: str, placeholder) -> str:
    """
    Stream the xAI API response and display it in real-time.
    Returns the complete response text.
    """
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {os.environ.get('xai')}"
    }
    
    data = {
        "messages": [
            {
                "role": "system",
                "content": "You are a helpful scientific research assistant. Analyze the following research query and provide initial insights."
            },
            {
                "role": "user",
                "content": query
            }
        ],
        "model": "grok-2-1212",
        "stream": True,
        "temperature": 0.7
    }
    
    try:
        response = requests.post(
            "https://api.x.ai/v1/chat/completions",
            headers=headers,
            json=data,
            stream=True,
            timeout=30
        )
        response.raise_for_status()
        
        full_response = ""
        
        for line in response.iter_lines():
            if line:
                line = line.decode('utf-8')
                if line.startswith('data: '):
                    json_str = line[6:]  # Remove 'data: ' prefix
                    if json_str == '[DONE]':
                        break
                    try:
                        chunk = json.loads(json_str)
                        if chunk["choices"][0]["delta"].get("content"):
                            content = chunk["choices"][0]["delta"]["content"]
                            full_response += content
                            # Update the placeholder with accumulated text
                            placeholder.markdown(full_response + "▌")
                    except json.JSONDecodeError:
                        continue
                        
        # Final update without the cursor
        placeholder.markdown(full_response)
        return full_response
        
    except requests.exceptions.RequestException as e:
        st.error(f"Error in streaming xAI API call: {str(e)}")
        return None

# -------------------- Main AI Lookup Function --------------------
def perform_ai_lookup(q, vocal_summary=True, extended_refs=False, titles_summary=True, full_audio=False, use_streaming=True):
    """Perform Arxiv search with initial xAI insights."""
    start = time.time()
    
    # First, get xAI insights
    st.write("### 🤖 Initial AI Insights")
    initial_insights = None
    
    if use_streaming:
        # Create a placeholder for streaming text
        streaming_placeholder = st.empty()
        with st.spinner("Getting streaming AI insights..."):
            initial_insights = stream_xai_response(q, streaming_placeholder)
    else:
        with st.spinner("Getting batch AI insights..."):
            xai_response = call_xai_api_batch(q)
            if xai_response and 'choices' in xai_response:
                initial_insights = xai_response['choices'][0]['message']['content']
                st.markdown(initial_insights)
    
    # Generate audio for xAI insights if enabled
    if vocal_summary and initial_insights:
        insights_text = clean_for_speech(initial_insights)
        if insights_text.strip():
            audio_file_insights = speak_with_edge_tts(insights_text)
            if audio_file_insights:
                st.write("### 🎤 AI Insights Audio")
                play_and_download_audio(audio_file_insights)
    
    # Proceed with existing ArXiv search
    st.write("### 📚 ArXiv Results")
    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")

    result = f"### 🔎 {q}\n\n{r2}\n\n{refs}"

    # Audio outputs
    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)
        if audio_file_full:
            st.write("### 📚 Full Audio")
            play_and_download_audio(audio_file_full)

    if vocal_summary:
        main_text = clean_for_speech(r2)
        if main_text.strip():
            audio_file_main = speak_with_edge_tts(main_text)
            if audio_file_main:
                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)
        if summaries_text.strip():
            audio_file_refs = speak_with_edge_tts(summaries_text)
            if audio_file_refs:
                st.write("### 📜 Long Refs")
                play_and_download_audio(audio_file_refs)

    if titles_summary:
        titles = []
        for line in 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)
            if titles_text.strip():
                audio_file_titles = speak_with_edge_tts(titles_text)
                if audio_file_titles:
                    st.write("### 🔖 Titles")
                    play_and_download_audio(audio_file_titles)

    st.markdown(result)
    
    # Save complete results including xAI insights
    if initial_insights:
        full_result = f"### 🤖 Initial AI Insights\n\n{initial_insights}\n\n{result}"
    else:
        full_result = result
        
    create_file(q, full_result, "md")
            
    elapsed = time.time()-start
    st.write(f"**Total Elapsed:** {elapsed:.2f} s")

    return full_result

# -------------------- Main Application --------------------
def main():
    st.session_state['user_name'] = st.selectbox("Current User:", USER_NAMES, index=0)

    # Display saved files in 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":
                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.button("🗑️ Clear All History in Sidebar"):
        md_files = glob.glob("*.md")
        mp3_files = glob.glob("*.mp3")
        for f in md_files+mp3_files:
            os.remove(f)
        st.success("All history cleared!")
        st.rerun()

    st.title("🎙️ ArXiv Voice Search")

    # Voice component
    mycomponent = components.declare_component("mycomponent", path="mycomponent")
    voice_val = mycomponent(my_input_value="Start speaking...")

    tabs = st.tabs(["🎤 Voice Chat", "💾 History", "⚙️ Settings"])

    with tabs[0]:
        st.subheader("🎤 Voice Chat")
        if voice_val:
            voice_text = voice_val.strip()
            input_changed = (voice_text != st.session_state.get('old_val'))
            if input_changed and voice_text:
                # Save user input
                create_file(st.session_state['user_name'], voice_text, "md")

                # Perform AI lookup with current streaming setting
                with st.spinner("Processing..."):
                    result = perform_ai_lookup(
                        voice_text,
                        vocal_summary=True,
                        extended_refs=False,
                        titles_summary=True,
                        full_audio=False,
                        use_streaming=st.session_state['use_streaming']
                    )

                st.session_state['old_val'] = voice_text

        st.write("Speak a query to run an ArXiv search and hear the results.")

    with tabs[1]:
        st.subheader("💾 History")
        # Show all MD files and allow reading them aloud
        md_files = sorted(glob.glob("*.md"), key=os.path.getmtime, reverse=True)
        for i, fpath in enumerate(md_files, start=1):
            fname = os.path.basename(fpath)
            with open(fpath,'r',encoding='utf-8') as ff:
                content = ff.read()
            with st.expander(fname, expanded=False):
                st.write(content)
                if st.button(f"🔊 Read Aloud {fname}", key=f"read_{i}_{fname}"):
                    voice = USER_VOICES.get(st.session_state['user_name'], "en-US-AriaNeural")
                    audio_file = speak_with_edge_tts(content, voice=voice)
                    if audio_file:
                        play_and_download_audio(audio_file)

        if st.button("📜 Read Entire History"):
            all_content = []
            for fpath in sorted(md_files, key=os.path.getmtime):
                with open(fpath,'r',encoding='utf-8') as ff:
                    c = ff.read().strip()
                    if c:
                        all_content.append((fpath, c))
            mp3_files = []
            for (fpath, text) in all_content:
                voice = USER_VOICES.get(st.session_state['user_name'], "en-US-AriaNeural")
                audio_file = speak_with_edge_tts(text, voice=voice)
                if audio_file:
                    mp3_files.append(audio_file)
                    st.write(f"**{os.path.basename(fpath)}:**")
                    play_and_download_audio(audio_file)

            if mp3_files:
                combined_file = f"full_conversation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp3"
                with open(combined_file, 'wb') as outfile:
                    for f in mp3_files:
                        with open(f, 'rb') as infile:
                            outfile.write(infile.read())
                st.write("**Full Conversation Audio:**")
                play_and_download_audio(combined_file)

    with tabs[2]:
        st.subheader("⚙️ Settings")
        st.session_state['use_streaming'] = st.toggle(
            "Use streaming responses",
            value=st.session_state['use_streaming'],
            help="Enable to see AI responses as they are generated in real-time"
        )

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

if __name__ == "__main__":
    main()