import PyPDF2
import docx
import os
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
import faiss
import pickle
import numpy as np
import streamlit as st
import io

def extract_text_from_pdf(file):
    pdf_reader = PyPDF2.PdfReader(io.BytesIO(file.read()))
    text = ""
    for page in pdf_reader.pages:
        text += page.extract_text() + "\n"
    return text

def extract_text_from_docx(file):
    doc = docx.Document(io.BytesIO(file.read()))
    text = ""
    for para in doc.paragraphs:
        text += para.text + "\n"
    return text

def extract_text_from_txt(file):
    return file.read().decode('utf-8')

def process_uploaded_files(uploaded_files):
    documents = []
    for uploaded_file in uploaded_files:
        filename = uploaded_file.name.lower()
        if filename.endswith('.pdf'):
            text = extract_text_from_pdf(uploaded_file)
        elif filename.endswith('.docx'):
            text = extract_text_from_docx(uploaded_file)
        elif filename.endswith('.txt'):
            text = extract_text_from_txt(uploaded_file)
        else:
            continue  # skip unsupported
        doc = Document(page_content=text, metadata={"source": uploaded_file.name})
        documents.append(doc)
    return documents

def create_vectorstore_from_docs(documents, embedding_model):
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    docs = text_splitter.split_documents(documents)
    
    embeddings = embedding_model.embed_documents([doc.page_content for doc in docs])
    embeddings = np.array(embeddings).astype('float32')
    
    dimension = embeddings.shape[1]
    index = faiss.IndexFlatIP(dimension)
    faiss.normalize_L2(embeddings)
    index.add(embeddings)
    
    # Save to temp
    temp_dir = "temp_vectorstore"
    os.makedirs(temp_dir, exist_ok=True)
    faiss.write_index(index, f"{temp_dir}/faiss.index")
    with open(f"{temp_dir}/documents.pkl", "wb") as f:
        pickle.dump(docs, f)
    
    return temp_dir

def retrieve_from_custom_db(query, db_path, embedding_model, top_k=10):
    formatted_query = f"query: {query.strip()}"
    query_vector = embedding_model.embed_query(formatted_query)
    query_vector = np.array(query_vector).astype('float32').reshape(1, -1)
    faiss.normalize_L2(query_vector)
    
    index_path = f"{db_path}/faiss.index"
    docstore_path = f"{db_path}/documents.pkl"
    
    index = faiss.read_index(index_path)
    with open(docstore_path, "rb") as f:
        documents = pickle.load(f)
    
    scores, indices = index.search(query_vector, top_k)
    results = []
    for score, idx in zip(scores[0], indices[0]):
        if idx != -1:
            doc = documents[idx]
            results.append((doc.page_content, score))
    return results