import os
import json
import re
import gradio as gr
import requests
from duckduckgo_search import DDGS
from typing import List
from pydantic import BaseModel, Field
from tempfile import NamedTemporaryFile
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.embeddings import HuggingFaceEmbeddings
from llama_parse import LlamaParse
from langchain_core.documents import Document
from huggingface_hub import InferenceClient
import inspect

# Environment variables and configurations
huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
llama_cloud_api_key = os.environ.get("LLAMA_CLOUD_API_KEY")

MODELS = [
    "google/gemma-2-9b",
    "mistralai/Mixtral-8x7B-Instruct-v0.1",
    "mistralai/Mistral-7B-Instruct-v0.3",
    "microsoft/Phi-3-mini-4k-instruct"
]

# Initialize LlamaParse
llama_parser = LlamaParse(
    api_key=llama_cloud_api_key,
    result_type="markdown",
    num_workers=4,
    verbose=True,
    language="en",
)

def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[Document]:
    """Loads and splits the document into pages."""
    if parser == "pypdf":
        loader = PyPDFLoader(file.name)
        return loader.load_and_split()
    elif parser == "llamaparse":
        try:
            documents = llama_parser.load_data(file.name)
            return [Document(page_content=doc.text, metadata={"source": file.name}) for doc in documents]
        except Exception as e:
            print(f"Error using Llama Parse: {str(e)}")
            print("Falling back to PyPDF parser")
            loader = PyPDFLoader(file.name)
            return loader.load_and_split()
    else:
        raise ValueError("Invalid parser specified. Use 'pypdf' or 'llamaparse'.")

def get_embeddings():
    return HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")

def update_vectors(files, parser):
    if not files:
        return "Please upload at least one PDF file."
    
    embed = get_embeddings()
    total_chunks = 0
    
    all_data = []
    for file in files:
        data = load_document(file, parser)
        all_data.extend(data)
        total_chunks += len(data)
    
    if os.path.exists("faiss_database"):
        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
        database.add_documents(all_data)
    else:
        database = FAISS.from_documents(all_data, embed)
    
    database.save_local("faiss_database")
    
    return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}."

def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=3, temperature=0.2, stop_clicked=None):
    print(f"Starting generate_chunked_response with {num_calls} calls")
    client = InferenceClient(model, token=huggingface_token)
    full_responses = []
    messages = [{"role": "user", "content": prompt}]
    
    for i in range(num_calls):
        print(f"Starting API call {i+1}")
        if (isinstance(stop_clicked, gr.State) and stop_clicked.value) or stop_clicked:
            print("Stop clicked, breaking loop")
            break
        try:
            response = ""
            for message in client.chat_completion(
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
                stream=True,
            ):
                if (isinstance(stop_clicked, gr.State) and stop_clicked.value) or stop_clicked:
                    print("Stop clicked during streaming, breaking")
                    break
                if message.choices and message.choices[0].delta and message.choices[0].delta.content:
                    chunk = message.choices[0].delta.content
                    response += chunk
            print(f"API call {i+1} response: {response[:100]}...")
            full_responses.append(response)
        except Exception as e:
            print(f"Error in generating response: {str(e)}")
    
    combined_response = " ".join(full_responses)
    print(f"Combined response: {combined_response[:100]}...")
    
    clean_response = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', combined_response, flags=re.DOTALL)
    clean_response = clean_response.replace("Using the following context:", "").strip()
    clean_response = clean_response.replace("Using the following context from the PDF documents:", "").strip()
    
    print(f"Final clean response: {clean_response[:100]}...")
    return clean_response

def duckduckgo_search(query):
    with DDGS() as ddgs:
        results = ddgs.text(query, max_results=5)
    return results

class CitingSources(BaseModel):
    sources: List[str] = Field(
        ...,
        description="List of sources to cite. Should be an URL of the source."
    )

def get_response_with_search(query, model, num_calls=3, temperature=0.2, stop_clicked=None):
    search_results = duckduckgo_search(query)
    context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n" 
                        for result in search_results if 'body' in result)
    
    prompt = f"""<s>[INST] Using the following context:
{context}
Write a detailed and complete research document that fulfills the following user request: '{query}'
After writing the document, please provide a list of sources used in your response. [/INST]"""
    
    generated_text = generate_chunked_response(prompt, model, num_calls=num_calls, temperature=temperature, stop_clicked=stop_clicked)
   
    # Clean the response
    clean_text = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', generated_text, flags=re.DOTALL)
    clean_text = clean_text.replace("Using the following context:", "").strip()
    
    # Split the content and sources
    parts = clean_text.split("Sources:", 1)
    main_content = parts[0].strip()
    sources = parts[1].strip() if len(parts) > 1 else ""
    
    return main_content, sources

def get_response_from_pdf(query, model, num_calls=3, temperature=0.2, stop_clicked=None):
    embed = get_embeddings()
    if os.path.exists("faiss_database"):
        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
    else:
        return "No documents available. Please upload PDF documents to answer questions."

    retriever = database.as_retriever()
    relevant_docs = retriever.get_relevant_documents(query)
    context_str = "\n".join([doc.page_content for doc in relevant_docs])

    prompt = f"""<s>[INST] Using the following context from the PDF documents:
{context_str}
Write a detailed and complete response that answers the following user question: '{query}'
Do not include a list of sources in your response. [/INST]"""

    generated_text = generate_chunked_response(prompt, model, num_calls=num_calls, temperature=temperature, stop_clicked=stop_clicked)

    # Clean the response
    clean_text = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', generated_text, flags=re.DOTALL)
    clean_text = clean_text.replace("Using the following context from the PDF documents:", "").strip()

    return clean_text

def chatbot_interface(message, history, use_web_search, model, temperature):
    if not message.strip():  # Check if the message is empty or just whitespace
        return history

    if use_web_search:
        main_content, sources = get_response_with_search(message, model, temperature)
        formatted_response = f"{main_content}\n\nSources:\n{sources}"
    else:
        response = get_response_from_pdf(message, model, temperature)
        formatted_response = response

    # Check if the last message in history is the same as the current message
    if history and history[-1][0] == message:
        # Replace the last response instead of adding a new one
        history[-1] = (message, formatted_response)
    else:
        # Add the new message-response pair
        history.append((message, formatted_response))

    return history


def clear_and_update_chat(message, history, use_web_search, model, temperature):
    updated_history = chatbot_interface(message, history, use_web_search, model, temperature)
    return "", updated_history  # Return empty string to clear the input

def retry_last_response(history):
    if history:
        last_user_message = history[-1][0]
        return last_user_message, history[:-1]
    return "", history

def undo_last_interaction(history):
    if len(history) >= 1:
        return history[:-1]
    return history

def clear_conversation():
    return []

def stop_generation():
    global is_generating
    is_generating = False

with gr.Blocks() as demo:
    is_generating = gr.State(False)
    stop_clicked = gr.State(False)

    gr.Markdown("# AI-powered Web Search and PDF Chat Assistant")
    
    with gr.Row():
        file_input = gr.Files(label="Upload your PDF documents", file_types=[".pdf"])
        parser_dropdown = gr.Dropdown(choices=["pypdf", "llamaparse"], label="Select PDF Parser", value="llamaparse")
        update_button = gr.Button("Upload Document")
    
    update_output = gr.Textbox(label="Update Status")
    update_button.click(update_vectors, inputs=[file_input, parser_dropdown], outputs=update_output)
    
    chatbot = gr.Chatbot(label="Conversation")
    msg = gr.Textbox(label="Ask a question")
    use_web_search = gr.Checkbox(label="Use Web Search", value=False)
    
    with gr.Row():
        model_dropdown = gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[1])
        temperature_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature")
        num_calls_slider = gr.Slider(minimum=1, maximum=5, value=3, step=1, label="Number of API Calls")
    
    with gr.Row():
        submit_btn = gr.Button("Send")
        stop_btn = gr.Button("Stop", visible=False)
        retry_btn = gr.Button("Retry")
        undo_btn = gr.Button("Undo")
        clear_btn = gr.Button("Clear")

    def protected_generate_response(message, history, use_web_search, model, temperature, num_calls, is_generating, stop_clicked):
        print("Starting protected_generate_response")
        if is_generating:
            print("Already generating, returning")
            return message, history, is_generating, stop_clicked
    
        is_generating = True
        
        if isinstance(stop_clicked, gr.State):
            stop_clicked.value = False
        else:
            stop_clicked = False
    
        try:
            print(f"Generating response for: {message}")
            if use_web_search:
                print("Using web search")
                main_content, sources = get_response_with_search(message, model, num_calls=num_calls, temperature=temperature, stop_clicked=stop_clicked)
                formatted_response = f"{main_content}\n\nSources:\n{sources}"
            else:
                print("Using PDF search")
                formatted_response = get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature, stop_clicked=stop_clicked)
    
            print(f"Generated response: {formatted_response[:100]}...")
    
        except Exception as e:
            print(f"Error generating response: {str(e)}")
            formatted_response = "I'm sorry, but I encountered an error while generating the response. Please try again."
        
        is_generating = False
        print(f"Returning final response")
        return "", history + [(message, formatted_response)], is_generating, stop_clicked

    def on_submit(message, history, use_web_search, model, temperature, num_calls, is_generating, stop_clicked):
        print(f"Submit button clicked with message: {message}")
        _, new_history, new_is_generating, new_stop_clicked = protected_generate_response(
            message, history, use_web_search, model, temperature, num_calls, is_generating, stop_clicked
        )
        print(f"New history has {len(new_history)} items")
        return "", new_history, new_is_generating, new_stop_clicked

    submit_btn.click(
        on_submit,
        inputs=[msg, chatbot, use_web_search, model_dropdown, temperature_slider, num_calls_slider, is_generating, stop_clicked],
        outputs=[msg, chatbot, is_generating, stop_clicked],
        show_progress=True
    ).then(
        lambda: gr.update(visible=True),
        None,
        stop_btn
    )

    stop_btn.click(
        lambda: (True, gr.update(visible=False)),
        None,
        [stop_clicked, stop_btn]
    )

    retry_btn.click(
        retry_last_response,
        inputs=[chatbot],
        outputs=[msg, chatbot]
    ).then(
        on_submit,
        inputs=[msg, chatbot, use_web_search, model_dropdown, temperature_slider, num_calls_slider, is_generating, stop_clicked],
        outputs=[msg, chatbot, is_generating, stop_clicked]
    )

    undo_btn.click(undo_last_interaction, inputs=[chatbot], outputs=[chatbot])
    clear_btn.click(clear_conversation, outputs=[chatbot])

    gr.Examples(
        examples=[
            ["What are the latest developments in AI?"],
            ["Tell me about recent updates on GitHub"],
            ["What are the best hotels in Galapagos, Ecuador?"],
            ["Summarize recent advancements in Python programming"],
        ],
        inputs=msg,
    )

    gr.Markdown(
    """
    ## How to use
    1. Upload PDF documents using the file input at the top.
    2. Select the PDF parser (pypdf or llamaparse) and click "Upload Document" to update the vector store.
    3. Ask questions in the textbox. 
    4. Toggle "Use Web Search" to switch between PDF chat and web search.
    5. Adjust Temperature and Number of API Calls sliders to fine-tune the response generation.
    6. Click "Send" or press Enter to get a response.
    7. Use "Retry" to regenerate the last response, "Undo" to remove the last interaction, and "Clear" to reset the conversation.
    8. Click "Stop" during generation to halt the process.
    """
    )
if __name__ == "__main__":
    demo.launch(share=True)