import torch
import transformers
import gradio as gr
from ragatouille import RAGPretrainedModel
from huggingface_hub import InferenceClient
import re
from datetime import datetime
import json
import arxiv
from utils import get_md_text_abstract, search_cleaner, get_arxiv_live_search
import os
import glob

# 🎛️ App configuration - tweak these knobs for maximum brain power! 🧠💪
retrieve_results = 20 
show_examples = True
llm_models_to_choose = ['mistralai/Mixtral-8x7B-Instruct-v0.1', 'mistralai/Mistral-7B-Instruct-v0.2', 'google/gemma-7b-it', 'None']

# 🎭 LLM acting instructions - "To be, or not to be... verbose" 🤔
generate_kwargs = dict(
    temperature = None,
    max_new_tokens = 512,
    top_p = None,
    do_sample = False,
)

# 🧙‍♂️ Summoning the RAG model - "Accio knowledge!" 📚✨
RAG = RAGPretrainedModel.from_index("colbert/indexes/arxiv_colbert")

try:
    gr.Info("🏗️ Setting up the knowledge retriever, please wait... 🕰️")
    rag_initial_output = RAG.search("What is Generative AI in Healthcare?", k = 1)
    gr.Info("🎉 Retriever is up and running! Time to flex those brain muscles! 💪🧠")
except:
    gr.Warning("😱 Oh no! The retriever took a coffee break. Try again later! ☕")

# 📜 The grand introduction - roll out the red carpet! 🎭
mark_text = '# 🩺🔍 Search Results\n'
header_text = "## 📚Arxiv📖Paper🔍Search - 🕵️‍♀️ Uncover, 📝 Summarize, and 🧩 Solve 🔬 Research 🤔❓ Puzzles ✍️ with 📚 Papers and 🤖 RAG AI 🧠\n"

# 🕰️ Time travel to find when our knowledge was last updated 🚀
try:
    with open("README.md", "r") as f:
        mdfile = f.read()
    date_pattern = r'Index Last Updated : \d{4}-\d{2}-\d{2}'
    match = re.search(date_pattern, mdfile)
    date = match.group().split(': ')[1]
    formatted_date = datetime.strptime(date, '%Y-%m-%d').strftime('%d %b %Y')
    header_text += f'Index Last Updated: {formatted_date}\n'
    index_info = f"Semantic Search - up to {formatted_date}"  
except:
    index_info = "Semantic Search"

database_choices = [index_info, 'Arxiv Search - Latest - (EXPERIMENTAL)']

# 🦉 Arxiv API - the wise old owl of academic knowledge 📜
arx_client = arxiv.Client()
is_arxiv_available = True
check_arxiv_result = get_arxiv_live_search("What is Self Rewarding AI and how can it be used in Multi-Agent Systems?", arx_client, retrieve_results)
if len(check_arxiv_result) == 0:
    is_arxiv_available = False
    print("😴 Arxiv search is taking a nap, switching to default search ...")
    database_choices = [index_info]

# 🎭 Show examples - a teaser trailer for your brain! 🍿🧠
sample_outputs = {
    'output_placeholder': 'The LLM will provide an answer to your question here...',
    'search_placeholder': '''
1. What is MoE?
2. What are Multi Agent Systems? 
3. What is Self Rewarding AI?
4. What is Semantic and Episodic memory?
5. What is AutoGen?
6. What is ChatDev?
7. What is Omniverse?
8. What is Lumiere?
9. What is SORA?
'''
}

output_placeholder = sample_outputs['output_placeholder']
md_text_initial = sample_outputs['search_placeholder']

# 🧹 Clean up the RAG output - nobody likes a messy mind! 🧼🧠
def rag_cleaner(inp):
    rank = inp['rank']
    title = inp['document_metadata']['title']
    content = inp['content']
    date = inp['document_metadata']['_time']
    return f"{rank}. <b> {title} </b> \n Date : {date} \n Abstract: {content}"

# 🎭 Craft the perfect prompt - it's showtime for the LLM! 🎬
def get_prompt_text(question, context, formatted = True, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2'):
    if formatted:
        sys_instruction = f"Context:\n {context} \n Given the following scientific paper abstracts, take a deep breath and let's think step by step to answer the question. Cite the titles of your sources when answering, do not cite links or dates."
        message = f"Question: {question}"
        
        if 'mistralai' in llm_model_picked:
            return f"<s>" + f"[INST] {sys_instruction}" +  f" {message}[/INST]"
        elif 'gemma' in llm_model_picked:
            return f"<bos><start_of_turn>user\n{sys_instruction}" +  f" {message}<end_of_turn>\n"
          
    return f"Context:\n {context} \n Given the following info, take a deep breath and let's think step by step to answer the question: {question}. Cite the titles of your sources when answering.\n\n"

# 🕵️‍♀️ Get those juicy references - time to go treasure hunting! 💎📚
def get_references(question, retriever, k = retrieve_results):
    rag_out = retriever.search(query=question, k=k)
    return rag_out

def get_rag(message):
    return get_references(message, RAG)

# 🎤 Save the response and read it aloud - it's karaoke time for your brain! 🧠🎶
def SaveResponseAndRead(result):
    documentHTML5='''
    <!DOCTYPE html>
    <html>
    <head>
        <title>Read It Aloud</title>
        <script type="text/javascript">
            function readAloud() {
                const text = document.getElementById("textArea").value;
                const speech = new SpeechSynthesisUtterance(text);
                window.speechSynthesis.speak(speech);
            }
        </script>
    </head>
    <body>
        <h1>🔊 Read It Aloud</h1>
        <textarea id="textArea" rows="10" cols="80">
    '''
    documentHTML5 = documentHTML5 + result
    documentHTML5 = documentHTML5 + '''
        </textarea>
        <br>
        <button onclick="readAloud()">🔊 Read Aloud</button>
    </body>
    </html>
    '''
    gr.HTML(documentHTML5)

# 📁 File management functions - because even AI needs a filing system! 🗄️🤖

def save_response_as_markdown(question, response):
    timestamp = datetime.now().strftime("%Y%m%d%H%M")
    filename = f"{timestamp}_{question[:50]}.md"  # Truncate question to 50 chars for filename
    with open(filename, "w", encoding="utf-8") as f:
        f.write(response)
    return filename

def list_markdown_files():
    files = glob.glob("*.md")
    files.sort(key=os.path.getmtime, reverse=True)
    return [f for f in files if f != "README.md"]

def delete_file(filename):
    if filename != "README.md":
        os.remove(filename)
        return f"Deleted {filename}"
    return "Cannot delete README.md"

def display_markdown_contents():
    files = list_markdown_files()
    output = ""
    for file in files:
        with open(file, "r", encoding="utf-8") as f:
            content = f.read()
        output += f"## {file}\n\n```markdown\n{content}\n```\n\n"
    return output

# 🎨 Building the UI - it's like LEGO, but for brains! 🧠🏗️
with gr.Blocks(theme = gr.themes.Soft()) as demo:
    header = gr.Markdown(header_text)
    
    with gr.Group():
        msg = gr.Textbox(label = 'Search', placeholder = 'What is Generative AI in Healthcare?')
        
        with gr.Accordion("Advanced Settings", open=False):
            with gr.Row(equal_height = True):
                llm_model = gr.Dropdown(choices = llm_models_to_choose, value = 'mistralai/Mistral-7B-Instruct-v0.2', label = 'LLM Model')
                llm_results = gr.Slider(minimum=4, maximum=10, value=5, step=1, interactive=True, label="Top n results as context")
                database_src = gr.Dropdown(choices = database_choices, value = index_info, label = 'Search Source')
                stream_results = gr.Checkbox(value = True, label = "Stream output", visible = False)

    output_text = gr.Textbox(show_label = True, container = True, label = 'LLM Answer', visible = True, placeholder = output_placeholder)
    input = gr.Textbox(show_label = False, visible = False)
    gr_md = gr.Markdown(mark_text + md_text_initial)

    with gr.Tab("Saved Responses"):
        refresh_button = gr.Button("🔄 Refresh File List")
        file_list = gr.Dropdown(choices=list_markdown_files(), label="Saved Responses")
        delete_button = gr.Button("🗑️ Delete Selected File")
        markdown_display = gr.Markdown()

    # 🔄 Update the file list - keeping things fresh! 🌿
    def update_file_list():
        return gr.Dropdown(choices=list_markdown_files())

    refresh_button.click(update_file_list, outputs=[file_list])
    delete_button.click(delete_file, inputs=[file_list], outputs=[markdown_display]).then(update_file_list, outputs=[file_list])
    file_list.change(lambda x: open(x, "r", encoding="utf-8").read() if x else "", inputs=[file_list], outputs=[markdown_display])

    # 🎭 The grand finale - where the magic happens! 🎩✨
    def update_with_rag_md(message, llm_results_use = 5, database_choice = index_info, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2'):
        prompt_text_from_data = ""
        database_to_use = database_choice
        if database_choice == index_info:
            rag_out = get_rag(message)
        else:
            arxiv_search_success = True
            try:
                rag_out = get_arxiv_live_search(message, arx_client, retrieve_results)
                if len(rag_out) == 0:
                    arxiv_search_success = False 
            except:
                arxiv_search_success = False

            if not arxiv_search_success:
                gr.Warning("😴 Arxiv Search is taking a siesta, switching to semantic search ...")
                rag_out = get_rag(message)
                database_to_use = index_info 

        md_text_updated = mark_text
        for i in range(retrieve_results):
            rag_answer = rag_out[i]
            if i < llm_results_use:
                md_text_paper, prompt_text = get_md_text_abstract(rag_answer, source = database_to_use, return_prompt_formatting = True)
                prompt_text_from_data += f"{i+1}. {prompt_text}"
            else:
                md_text_paper = get_md_text_abstract(rag_answer, source = database_to_use)
            md_text_updated += md_text_paper
        prompt = get_prompt_text(message, prompt_text_from_data, llm_model_picked = llm_model_picked)
        return md_text_updated, prompt

    # 🧠 Asking the LLM - it's like a really smart magic 8-ball! 🎱✨
    def ask_llm(prompt, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2', stream_outputs = False):
        model_disabled_text = "LLM Model is taking a vacation. Try again later! 🏖️"
        output = ""
        
        if llm_model_picked == 'None':
            if stream_outputs:
                for out in model_disabled_text:
                    output += out
                    yield output
                return output 
            else:
                return model_disabled_text
              
        client = InferenceClient(llm_model_picked)
        try:
            stream = client.text_generation(prompt, **generate_kwargs, stream=stream_outputs, details=False, return_full_text=False)
           
        except:
            gr.Warning("🚦 LLM Inference hit a traffic jam! Take a breather and try again later.")
            return ""
       
        if stream_outputs:
            for response in stream:
                output += response
                SaveResponseAndRead(response)
                yield output
            return output
        else:
            return stream

    # 🎬 Action! Process the query and save the response
    def process_and_save(message, llm_results_use, database_choice, llm_model_picked):
        md_text_updated, prompt = update_with_rag_md(message, llm_results_use, database_choice, llm_model_picked)
        llm_response = ask_llm(prompt, llm_model_picked, stream_outputs=False)
        full_response = f"Question: {message}\n\nResponse:\n{llm_response}\n\nReferences:\n{md_text_updated}"
        filename = save_response_as_markdown(message, full_response)
        return md_text_updated, prompt, llm_response, filename

    # 🎬 Lights, camera, action! Let's get this show on the road! 🚀
    msg.submit(process_and_save, [msg, llm_results, database_src, llm_model], [gr_md, input, output_text, file_list]).then(update_file_list, outputs=[file_list])

# 🎉 Launch the app - let the knowledge party begin! 🎊🧠
demo.queue().launch()