import os
import json
import re
import gradio as gr
import pandas as pd
import requests
import random
import urllib.parse
from tempfile import NamedTemporaryFile
from typing import List
from bs4 import BeautifulSoup
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import PyPDFLoader
from langchain_core.output_parsers import StrOutputParser
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.llms import HuggingFaceHub
from langchain_core.runnables import RunnableParallel, RunnablePassthrough
from langchain_core.documents import Document
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")

# Memory database to store question-answer pairs
memory_database = {}
conversation_history = []

def load_and_split_document_basic(file):
    """Loads and splits the document into pages."""
    loader = PyPDFLoader(file.name)
    data = loader.load_and_split()
    return data

def load_and_split_document_recursive(file: NamedTemporaryFile) -> List[Document]:
    """Loads and splits the document into chunks."""
    loader = PyPDFLoader(file.name)
    pages = loader.load()
    
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len,
    )
    
    chunks = text_splitter.split_documents(pages)
    return chunks

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

def create_or_update_database(data, embeddings):
    if os.path.exists("faiss_database"):
        db = FAISS.load_local("faiss_database", embeddings, allow_dangerous_deserialization=True)
        db.add_documents(data)
    else:
        db = FAISS.from_documents(data, embeddings)
    db.save_local("faiss_database")

def clear_cache():
    if os.path.exists("faiss_database"):
        os.remove("faiss_database")
        return "Cache cleared successfully."
    else:
        return "No cache to clear."

def get_similarity(text1, text2):
    vectorizer = TfidfVectorizer().fit_transform([text1, text2])
    return cosine_similarity(vectorizer[0:1], vectorizer[1:2])[0][0]

prompt = """
Answer the question based on the following information:

Conversation History:
{history}

Context from documents:
{context}

Current Question: {question}

If the question is referring to the conversation history, use that information to answer.
If the question is not related to the conversation history, use the context from documents to answer.
If you don't have enough information to answer, say so.

Provide a concise and direct answer to the question:
"""

def get_model(temperature, top_p, repetition_penalty):
    return HuggingFaceHub(
        repo_id="mistralai/Mistral-7B-Instruct-v0.3",
        model_kwargs={
            "temperature": temperature,
            "top_p": top_p,
            "repetition_penalty": repetition_penalty,
            "max_length": 1000
        },
        huggingfacehub_api_token=huggingface_token
    )

def generate_chunked_response(model, prompt, max_tokens=1000, max_chunks=5):
    full_response = ""
    for i in range(max_chunks):
        chunk = model(prompt + full_response, max_new_tokens=max_tokens)
        chunk = chunk.strip()
        if chunk.endswith((".", "!", "?")):
            full_response += chunk
            break
        full_response += chunk
    return full_response.strip()

def manage_conversation_history(question, answer, history, max_history=5):
    history.append({"question": question, "answer": answer})
    if len(history) > max_history:
        history.pop(0)
    return history

def is_related_to_history(question, history, threshold=0.3):
    if not history:
        return False
    history_text = " ".join([f"{h['question']} {h['answer']}" for h in history])
    similarity = get_similarity(question, history_text)
    return similarity > threshold

def extract_text_from_webpage(html):
    soup = BeautifulSoup(html, 'html.parser')
    for script in soup(["script", "style"]):
        script.extract()  # Remove scripts and styles
    text = soup.get_text()
    lines = (line.strip() for line in text.splitlines())
    chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
    text = '\n'.join(chunk for chunk in chunks if chunk)
    return text

_useragent_list = [
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Edge/91.0.864.59 Safari/537.36",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Edge/91.0.864.59 Safari/537.36",
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Safari/537.36",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Safari/537.36",
]

def google_search(term, num_results=5, lang="en", timeout=5, safe="active", ssl_verify=None):
    escaped_term = urllib.parse.quote_plus(term)
    start = 0
    all_results = []
    max_chars_per_page = 8000  # Limit the number of characters from each webpage to stay under the token limit

    print(f"Starting Google search for term: '{term}'")

    with requests.Session() as session:
        while start < num_results:
            try:
                user_agent = random.choice(_useragent_list)
                headers = {
                    'User-Agent': user_agent
                }
                resp = session.get(
                    url="https://www.google.com/search",
                    headers=headers,
                    params={
                        "q": term,
                        "num": num_results - start,
                        "hl": lang,
                        "start": start,
                        "safe": safe,
                    },
                    timeout=timeout,
                    verify=ssl_verify,
                )
                resp.raise_for_status()
                print(f"Successfully retrieved search results page (start={start})")
            except requests.exceptions.RequestException as e:
                print(f"Error retrieving search results: {e}")
                break

            soup = BeautifulSoup(resp.text, "html.parser")
            result_block = soup.find_all("div", attrs={"class": "g"})
            if not result_block:
                print("No results found on this page")
                break
            
            print(f"Found {len(result_block)} results on this page")
            for result in result_block:
                link = result.find("a", href=True)
                if link:
                    link = link["href"]
                    print(f"Processing link: {link}")
                    try:
                        webpage = session.get(link, headers=headers, timeout=timeout)
                        webpage.raise_for_status()
                        visible_text = extract_text_from_webpage(webpage.text)
                        if len(visible_text) > max_chars_per_page:
                            visible_text = visible_text[:max_chars_per_page] + "..."
                        all_results.append({"link": link, "text": visible_text})
                        print(f"Successfully extracted text from {link}")
                    except requests.exceptions.RequestException as e:
                        print(f"Error retrieving webpage content: {e}")
                        all_results.append({"link": link, "text": None})
                else:
                    print("No link found for this result")
                    all_results.append({"link": None, "text": None})
            start += len(result_block)

    print(f"Search completed. Total results: {len(all_results)}")
    print("Search results:")
    for i, result in enumerate(all_results, 1):
        print(f"Result {i}:")
        print(f"  Link: {result['link']}")
        if result['text']:
            print(f"  Text: {result['text'][:100]}...")  # Display the first 100 characters of the text for brevity
        else:
            print("  No text extracted")
    return all_results

def process_question(question, documents, history, temperature, top_p, repetition_penalty, enable_web_search):
    global conversation_history

    embeddings = get_embeddings()

    # Check the memory database for similar questions
    for prev_question, prev_answer in memory_database.items():
        similarity = get_similarity(question, prev_question)
        if similarity > 0.8:
            return prev_answer

    # Retrieve relevant documents from the vector store
    if os.path.exists("faiss_database"):
        db = FAISS.load_local("faiss_database", embeddings, allow_dangerous_deserialization=True)
        relevant_docs = db.similarity_search(question, k=3)
    else:
        relevant_docs = []

    # Perform web search if enabled and no relevant documents found
    if enable_web_search and len(relevant_docs) == 0:
        web_search_results = google_search(question, num_results=5)
        web_docs = [Document(page_content=res["text"] or "", metadata={"source": res["link"]}) for res in web_search_results if res["text"]]

        if web_docs:
            # Update the FAISS vector store with new documents
            create_or_update_database(web_docs, embeddings)
            db = FAISS.load_local("faiss_database", embeddings, allow_dangerous_deserialization=True)
            relevant_docs = db.similarity_search(question, k=3)

    context = "\n\n".join([doc.page_content for doc in relevant_docs])

    if is_related_to_history(question, history):
        context = "None"
    else:
        history_text = "\n".join([f"Q: {h['question']}\nA: {h['answer']}" for h in history]) if history else "None"
        context = context if context else "None"

    prompt_text = ChatPromptTemplate(
        input_variables=["history", "context", "question"],
        template=prompt
    ).format(history=history_text, context=context, question=question)

    model = get_model(temperature, top_p, repetition_penalty)
    answer = generate_chunked_response(model, prompt_text)
    
    conversation_history = manage_conversation_history(question, answer, history)
    memory_database[question] = answer

    return answer

def process_uploaded_file(file, is_recursive):
    if is_recursive:
        data = load_and_split_document_recursive(file)
    else:
        data = load_and_split_document_basic(file)
    embeddings = get_embeddings()
    create_or_update_database(data, embeddings)
    return "File processed and data added to the vector database."

def extract_db_to_excel():
    embed = get_embeddings()
    database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
    
    documents = database.docstore._dict.values()
    data = [{"page_content": doc.page_content, "metadata": json.dumps(doc.metadata)} for doc in documents]
    df = pd.DataFrame(data)
    
    with NamedTemporaryFile(delete=False, suffix='.xlsx') as tmp:
        excel_path = tmp.name
        df.to_excel(excel_path, index=False)
    
    return excel_path

def export_memory_db_to_excel():
    data = [{"question": question, "answer": answer} for question, answer in memory_database.items()]
    df_memory = pd.DataFrame(data)
    
    data_history = [{"question": item["question"], "answer": item["answer"]} for item in conversation_history]
    df_history = pd.DataFrame(data_history)
    
    with NamedTemporaryFile(delete=False, suffix='.xlsx') as tmp:
        excel_path = tmp.name
        with pd.ExcelWriter(excel_path, engine='openpyxl') as writer:
            df_memory.to_excel(writer, sheet_name='Memory Database', index=False)
            df_history.to_excel(writer, sheet_name='Conversation History', index=False)
    
    return excel_path

with gr.Blocks() as demo:
    with gr.Row():
        pdf_file = gr.File(label="Upload PDF")
    with gr.Row():
        recursive_check = gr.Checkbox(label="Use Recursive Text Splitter")
        upload_button = gr.Button("Upload and Process")
    with gr.Row():
        upload_output = gr.Textbox(label="Upload Output")
    
    with gr.Row():
        question = gr.Textbox(label="Your Question")
    with gr.Row():
        temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, label="Temperature")
        top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.9, label="Top P")
        repetition_penalty = gr.Slider(minimum=0.0, maximum=2.0, value=1.0, label="Repetition Penalty")
        web_search_check = gr.Checkbox(label="Enable Web Search")
    with gr.Row():
        ask_button = gr.Button("Ask")
    with gr.Row():
        answer = gr.Textbox(label="Answer")
    
    with gr.Row():
        clear_button = gr.Button("Clear Cache")
    with gr.Row():
        clear_output = gr.Textbox(label="Clear Output")

    with gr.Row():
        export_db_button = gr.Button("Export Database to Excel")
        export_db_output = gr.Textbox(label="Export Output")
    with gr.Row():
        export_memory_button = gr.Button("Export Memory DB to Excel")
        export_memory_output = gr.Textbox(label="Export Output")
    
    upload_button.click(process_uploaded_file, [pdf_file, recursive_check], upload_output)
    ask_button.click(process_question, [question, pdf_file, conversation_history, temperature, top_p, repetition_penalty, web_search_check], answer)
    clear_button.click(clear_cache, [], clear_output)
    export_db_button.click(extract_db_to_excel, [], export_db_output)
    export_memory_button.click(export_memory_db_to_excel, [], export_memory_output)

demo.launch()