import os
import gradio as gr
import logging
from groq import Groq
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import PyPDF2
from sklearn.metrics.pairwise import cosine_similarity
from collections import Counter

# --------------------- Setup ---------------------

logging.basicConfig(
    filename='query_logs.log',
    level=logging.INFO,
    format='%(asctime)s:%(levelname)s:%(message)s'
)

GROQ_API_KEY = "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"
client = Groq(api_key=GROQ_API_KEY)
PDF_PATH = 'Robert Ciesla - The Book of Chatbots_ From ELIZA to ChatGPT-Springer (2024).pdf'
sentence_transformer_model = SentenceTransformer('all-MiniLM-L6-v2')
cache = {}

# --------------------- Vectorization Function ---------------------

def vectorize_text(sentences_with_pages):
    """Vectorize sentences using SentenceTransformer and create a FAISS index."""
    try:
        sentences = [item['sentence'] for item in sentences_with_pages]
        embeddings = sentence_transformer_model.encode(sentences, show_progress_bar=True)
        index = faiss.IndexFlatL2(embeddings.shape[1])
        index.add(np.array(embeddings))
        logging.info(f"Added {len(sentences)} sentences to the vector store.")
        return index, sentences_with_pages
    except Exception as e:
        logging.error(f"Error during vectorization: {str(e)}")
        return None, None

# --------------------- PDF Processing ---------------------

def read_pdf(file_path):
    if not os.path.exists(file_path):
        logging.error(f"PDF file not found at: {file_path}")
        return []

    sentences_with_pages = []
    with open(file_path, 'rb') as file:
        reader = PyPDF2.PdfReader(file)
        for page_num, page in enumerate(reader.pages):
            text = page.extract_text()
            if text:
                sentences = [sentence.strip() for sentence in text.split('\n') if sentence.strip()]
                for sentence in sentences:
                    sentences_with_pages.append({'sentence': sentence, 'page_number': page_num + 1})
    return sentences_with_pages

# Read and Vectorize PDF Content
sentences_with_pages = read_pdf(PDF_PATH)
vector_index, sentences_with_pages = vectorize_text(sentences_with_pages)

# --------------------- Query Handling ---------------------

def generate_query_embedding(query):
    return sentence_transformer_model.encode([query])

def is_query_relevant(distances, threshold=1.0):
    return distances[0][0] <= threshold

def generate_diverse_responses(prompt, n=3):
    responses = []
    for i in range(n):
        temperature = 0.7 + (i * 0.1)
        top_p = 0.9 - (i * 0.1)
        try:
            chat_completion = client.chat.completions.create(
                messages=[{"role": "user", "content": prompt}],
                model="llama3-8b-8192",
                temperature=temperature,
                top_p=top_p
            )
            responses.append(chat_completion.choices[0].message.content.strip())
        except Exception as e:
            logging.error(f"Error generating response: {str(e)}")
            responses.append("Error generating this response.")
    return responses

def aggregate_responses(responses):
    response_counter = Counter(responses)
    most_common_response, count = response_counter.most_common(1)[0]
    if count > 1:
        return most_common_response
    else:
        embeddings = sentence_transformer_model.encode(responses)
        avg_embedding = np.mean(embeddings, axis=0)
        similarities = cosine_similarity([avg_embedding], embeddings)[0]
        return responses[np.argmax(similarities)]

def generate_answer(query):
    if query in cache:
        logging.info(f"Cache hit for query: {query}")
        return cache[query]

    try:
        query_embedding = generate_query_embedding(query)
        D, I = vector_index.search(np.array(query_embedding), k=5)

        if is_query_relevant(D):
            relevant_items = [sentences_with_pages[i] for i in I[0]]
            combined_text = " ".join([item['sentence'] for item in relevant_items])
            page_numbers = sorted(set([item['page_number'] for item in relevant_items]))
            page_numbers_str = ', '.join(map(str, page_numbers))

            # Construct primary prompt
            prompt = f"""
Use the following context from "The Book of Chatbots" to answer the question. If additional explanation is needed, provide an example.
**Context (Pages {page_numbers_str}):**
{combined_text}
**User's question:**
{query}
**Remember to indicate the specific page numbers.**
"""
            primary_responses = generate_diverse_responses(prompt)
            primary_answer = aggregate_responses(primary_responses)
            
            # Construct additional prompt for explanations
            explanation_prompt = f"""
The user has a question about a complex topic. Could you provide an explanation or example and real-life example for better understanding?
**User's question:**
{query}
**Primary answer:**
{primary_answer}
"""
            explanation_responses = generate_diverse_responses(explanation_prompt)
            explanation_answer = aggregate_responses(explanation_responses)

            # Combine primary answer and explanation
            full_response = f"{primary_answer}\n\n{explanation_answer}\n\n_From 'The Book of Chatbots,' pages {page_numbers_str}_"
            cache[query] = full_response
            logging.info(f"Generated response for query: {query}")
            return full_response

        else:
            # General knowledge fallback
            prompt = f"""
The user asked a question that is not covered in "The Book of Chatbots." Please provide a helpful answer using general knowledge.
**User's question:**
{query}
"""
            fallback_responses = generate_diverse_responses(prompt)
            fallback_answer = aggregate_responses(fallback_responses)
            cache[query] = fallback_answer
            return fallback_answer

    except Exception as e:
        logging.error(f"Error generating answer: {str(e)}")
        return "Sorry, an error occurred while generating the answer."

# --------------------- Gradio Interface ---------------------

def gradio_interface(user_query, history):
    response = generate_answer(user_query)
    history = history or []
    history.append({"role": "user", "content": user_query})
    history.append({"role": "assistant", "content": response})
    return history, history

# Create the Gradio interface
with gr.Blocks(css=".gradio-container {background-color: #f0f0f0}") as iface:
    gr.Markdown("""
    # **The Book of Chatbot**
    *Dive into the evolution of chatbots from ELIZA to ChatGPT with Chatbot Chronicles. Ask any question and explore the fascinating world of conversational AI as presented in Robert Ciesla's "The Book of Chatbots.*
    """)

    chatbot = gr.Chatbot(height=500, type='messages')
    state = gr.State([])

    with gr.Row():
        txt = gr.Textbox(
            show_label=False,
            placeholder="Type your message here and press Enter",
            container=False
        )
        submit_btn = gr.Button("Send")

    def submit_message(user_query, history):
        history = history or []
        history.append({"role": "user", "content": user_query})
        return "", history

    def bot_response(history):
        user_query = history[-1]['content']
        response = generate_answer(user_query)
        history.append({"role": "assistant", "content": response})
        return history

    txt.submit(submit_message, [txt, state], [txt, state], queue=False).then(
        bot_response, state, chatbot
    )
    submit_btn.click(submit_message, [txt, state], [txt, state], queue=False).then(
        bot_response, state, chatbot
    )

    reset_btn = gr.Button("Reset Chat")
    reset_btn.click(lambda: ([], []), outputs=[chatbot, state], queue=False)

# Launch the Gradio app
if __name__ == "__main__":
    iface.launch()