import torch
import fitz  # PyMuPDF for PDF text extraction
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from sentence_transformers import SentenceTransformer
import faiss
import gradio as gr
import os
from huggingface_hub import login

# Authenticate with Hugging Face Hub

pdf_path1 ='/content/Chrono 1.pdf'
pdf_path2 ='/content/Chrono 2.pdf'
pdf_path3 ='/content/Chrono 3.pdf'

# Load the Mistral 7B model and tokenizer
model_name = 'mistralai/Mistral-7B-Instruct-v0.3'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)

# Load sentence transformer for embedding and similarity search
embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

# Function to extract text from PDFs
def extract_text_from_pdf(pdf_file_path):
    doc = fitz.open(pdf_file_path)
    text = ""
    for page in doc:
        text += page.get_text("text")
    return text

# Placeholder PDF knowledge base (Extracted content from PDFs)
pdf_knowledge_base = []
pdf_files = [pdf_path1, pdf_path2,pdf_path3]  # Add the actual paths to your PDF files

for pdf_file in pdf_files:
    pdf_text = extract_text_from_pdf(pdf_file)
    pdf_knowledge_base.append({"document": pdf_file, "content": pdf_text})

# Combine extracted text with specific company information
knowledge_base = [
    {"question": "How does DiabeTrek ensure data privacy and security?",
     "answer": ("DiabeTrek ensures data privacy through multiple layers of protection, including data encryption during "
                "transit and at rest. We comply with regulations like HIPAA and GDPR to safeguard your personal data.")},
    {"question": "What are DiabeTrek's emergency guidelines?",
     "answer": "DiabeTrek advises users to seek immediate medical attention in case of diabetes-related emergencies. This chatbot is not for emergency use."},
    {"question": "What are DiabeTrek's mission, vision, and values?",
     "answer": "DiabeTrek's mission is to improve the lives of people with diabetes through innovative AI-driven solutions. Our vision is a world where diabetes care is seamless, proactive, and accessible."},
    # Additional items can be added here following the CEO's instructions
]

# Create a FAISS index for efficient retrieval
embedding_dim = 384  # Output dimension of the MiniLM model
index = faiss.IndexFlatL2(embedding_dim)

# Create a list of embeddings and index them
knowledge_embeddings = []
for entry in knowledge_base:
    embedding = embedder.encode(entry['question'], convert_to_tensor=False)
    knowledge_embeddings.append(embedding)
    index.add(embedding.reshape(1, -1))

# Create embeddings for PDF content and index them
for pdf_entry in pdf_knowledge_base:
    embedding = embedder.encode(pdf_entry['content'], convert_to_tensor=False)
    knowledge_embeddings.append(embedding)
    index.add(embedding.reshape(1, -1))

# RAG Retrieval function
def retrieve_knowledge(question, top_k=1):
    question_embedding = embedder.encode(question, convert_to_tensor=False)
    D, I = index.search(question_embedding.reshape(1, -1), top_k)
    results = [knowledge_base[idx] for idx in I[0]]
    return results

# Chatbot function combining retrieval and generation
def customer_support_chatbot(user_input):
    # Retrieve relevant knowledge
    retrieved_knowledge = retrieve_knowledge(user_input)

    # Prepare context for the generative model
    context = " ".join([f"Q: {entry['question']} A: {entry['answer']}" for entry in retrieved_knowledge])

    # Generate response using Mistral
    prompt = f"Customer Question: {user_input}\n\n{context}\n\nResponse:"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_length=150, do_sample=True, temperature=0.7)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return response

# Gradio UI
def gradio_interface(user_input):
    response = customer_support_chatbot(user_input)
    return response

# Build Gradio interface
interface = gr.Interface(fn=gradio_interface,
                         inputs="text",
                         outputs="text",
                         title="DiabeTrek Customer Support Chatbot",
                         description="Ask any question about DiabeTrek, its services, and policies.")

# Launch the Gradio app
interface.launch()