import gradio as gr
import pandas as pd
import numpy as np
import os
import re
import time
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import faiss
import docx2txt
import textract
from openpyxl import load_workbook

# File reading functions
def read_file(file_path):
    ext = os.path.splitext(file_path)[1].lower()
    try:
        if ext == '.txt':
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                return f.read()
        elif ext == '.docx':
            return docx2txt.process(file_path)
        elif ext == '.doc':
            return textract.process(file_path).decode('utf-8')
        elif ext in ('.xls', '.xlsx'):
            wb = load_workbook(file_path)
            return '\n'.join([str(cell.value) for sheet in wb for row in sheet for cell in row])
        elif ext == '.csv':
            return pd.read_csv(file_path).to_string()
    except Exception as e:
        print(f"Error reading {file_path}: {str(e)}")
    return ""

# Data normalization
def normalize_content(content):
    content = re.sub(r'\s+', ' ', content)  # Remove extra whitespace
    content = re.sub(r'[^\w\s.,;:!?\-]', '', content)  # Remove special chars
    return content[:10000]  # Limit length

# Create knowledge base
def create_knowledge_base():
    knowledge = []
    for root, _, files in os.walk("data"):
        for file in files:
            if file.split('.')[-1].lower() in ['txt', 'doc', 'docx', 'xls', 'xlsx', 'csv']:
                file_path = os.path.join(root, file)
                content = read_file(file_path)
                if content:
                    knowledge.append({
                        "file": file,
                        "path": file_path,
                        "content": normalize_content(content)
                    })
    return pd.DataFrame(knowledge)

# Build search index
def build_search_index(df):
    if df.empty:
        return None, None, None
    
    vectorizer = TfidfVectorizer(stop_words='english', max_features=10000)
    tfidf_matrix = vectorizer.fit_transform(df['content'])
    
    # Convert to float32 for FAISS
    vectors = tfidf_matrix.astype(np.float32).toarray()
    faiss.normalize_L2(vectors)
    
    # Create FAISS index
    dimension = vectors.shape[1]
    index = faiss.IndexFlatIP(dimension)
    index.add(vectors)
    
    return vectorizer, index, df

# Search function
def search(query, vectorizer, index, df, top_k=3):
    if vectorizer is None or index is None:
        return "No documents indexed yet!"
    
    query_vec = vectorizer.transform([query]).astype(np.float32).toarray()
    faiss.normalize_L2(query_vec)
    
    distances, indices = index.search(query_vec, top_k)
    
    results = []
    for i, (score, idx) in enumerate(zip(distances[0], indices[0])):
        if score > 0.1:  # Relevance threshold
            doc = df.iloc[idx]
            results.append(
                f"📄 **{doc['file']}** (Relevance: {score:.2f})\n"
                f"```\n{doc['content'][:300]}...\n```"
            )
    
    return "\n\n".join(results) if results else "No relevant documents found"

# Initialize system
print("Initializing knowledge base...")
start_time = time.time()
df = create_knowledge_base()
vectorizer, index, df = build_search_index(df)
load_time = time.time() - start_time
print(f"Indexed {len(df)} documents in {load_time:.1f} seconds")

# Gradio interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🔍 Knowledge Base Chatbot")
    gr.Markdown(f"**Indexed {len(df)} documents** | *Ask anything about your knowledge base*")
    
    with gr.Row():
        query = gr.Textbox(label="Search Query", placeholder="Ask about your documents...", scale=5)
        search_btn = gr.Button("Search", variant="primary")
    
    results = gr.Markdown()
    
    gr.Examples(
        examples=["What are the key metrics?", "Show me product specifications", "Explain the onboarding process"],
        inputs=query
    )
    
    gr.Markdown("### Supported file types: TXT, DOC, DOCX, XLS, XLSX, CSV")
    
    def perform_search(query):
        return search(query, vectorizer, index, df)
    
    query.submit(perform_search, inputs=query, outputs=results)
    search_btn.click(perform_search, inputs=query, outputs=results)

demo.launch()