import os
import gradio as gr
from transformers import AutoTokenizer, AutoModel
import torch
from transformers import pipeline
import wikipedia
import numpy as np
import faiss

# Initialize components
tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
vector_store = None

# Mean Pooling - Take attention mask into account for correct averaging
def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0]
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

def encode_text(texts):
    # Tokenize sentences
    encoded_input = tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors='pt')
    
    # Compute token embeddings
    with torch.no_grad():
        model_output = model(**encoded_input)

    # Perform pooling
    sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
    
    return sentence_embeddings.numpy()

class VectorStore:
    def __init__(self, dimension: int):
        self.dimension = dimension
        self.index = faiss.IndexFlatL2(dimension)
        self.texts = []

    def add_texts(self, texts, embeddings):
        self.texts.extend(texts)
        self.index.add(embeddings)

    def search(self, query_embedding, k):
        query_embedding = query_embedding.reshape(1, -1)
        distances, indices = self.index.search(query_embedding, k)
        results = []
        for idx, distance in zip(indices[0], distances[0]):
            if idx < len(self.texts):
                results.append((self.texts[idx], float(distance)))
        return results

def get_wikipedia_content():
    try:
        topics = [
            'Artificial intelligence',
            'Machine learning',
            'Deep learning',
            'Natural language processing',
            'Computer vision'
        ]
        
        all_content = []
        for topic in topics:
            try:
                page = wikipedia.page(topic)
                all_content.append(page.content[:2000])
            except wikipedia.exceptions.DisambiguationError as e:
                try:
                    page = wikipedia.page(e.options[0])
                    all_content.append(page.content[:2000])
                except:
                    continue
            except:
                continue
                
        return '\n\n'.join(all_content)
    except Exception as e:
        return """
        Artificial Intelligence (AI) is the simulation of human intelligence processes by machines.
        Machine Learning is a subset of artificial intelligence that provides systems the ability to learn.
        Deep Learning is part of machine learning based on artificial neural networks.
        Natural Language Processing (NLP) helps computers understand human language.
        """

def chunk_text(text, chunk_size=300):
    sentences = [s.strip() for s in text.split('.') if s.strip()]
    chunks = []
    current_chunk = []
    current_length = 0
    
    for sentence in sentences:
        sentence_length = len(sentence)
        if current_length + sentence_length > chunk_size and current_chunk:
            chunks.append('. '.join(current_chunk) + '.')
            current_chunk = [sentence]
            current_length = sentence_length
        else:
            current_chunk.append(sentence)
            current_length += sentence_length
    
    if current_chunk:
        chunks.append('. '.join(current_chunk) + '.')
    
    return chunks

def init_vector_store():
    global vector_store
    text = get_wikipedia_content()
    chunks = chunk_text(text)
    embeddings = encode_text(chunks)
    vector_store = VectorStore(dimension=embeddings.shape[1])
    vector_store.add_texts(chunks, embeddings)
    return vector_store

def generate_answer(query: str, context: str) -> str:
    try:
        generator = pipeline('text-generation', model='gpt2')
        prompt = f"Context: {prompt}"
        response = generator(prompt, max_length=150, num_return_sequences=1)
        return response[0]['generated_text']
    except Exception as e:
        relevant_sentences = [s for s in context.split('.') if query.lower() in s.lower()]
        if relevant_sentences:
            return relevant_sentences[0] + '.'
        return "I apologize, but I couldn't generate a specific answer based on the available information."

def chat_function(message, history):
    query_embedding = encode_text([message])
    results = vector_store.search(query_embedding, k=3)
    context = ' '.join([text for text, _ in results])
    answer = generate_answer(message, context)
    return answer

def handle_file_upload(file):
    if file is None:
        return "No file uploaded"
    
    try:
        with open(file.name, 'r', encoding='utf-8') as f:
            content = f.read()
        chunks = chunk_text(content)
        embeddings = encode_text(chunks)
        vector_store.add_texts(chunks, embeddings)
        return f"File processed successfully. Added {len(chunks)} chunks to knowledge base."
    except Exception as e:
        return f"Error processing file: {str(e)}"

# Initialize vector store
print("Initializing vector store with Wikipedia content...")
init_vector_store()

# Create Gradio interface
demo = gr.Blocks()

with demo:
    gr.Markdown("# RAG Chatbot")
    
    with gr.Row():
        file_input = gr.File(label="Upload Document")
        upload_button = gr.Button("Process File")
    
    upload_output = gr.Textbox(label="Upload Status")
    
    chatbot = gr.ChatInterface(
        chat_function,
        examples=["What is artificial intelligence?", "Explain machine learning", "Deep learning","Natural language processing","Computer vision"],
        title="Chat with your documents"
    )
    
    upload_button.click(
        handle_file_upload,
        inputs=[file_input],
        outputs=[upload_output]
    )

# Launch for Hugging Face Spaces
if __name__ == "__main__":
    demo.launch()