# =======================================
# 📘 RAG App Pro – Gemini + Smart Embeddings (Multi-User Safe)
# =======================================

import os, re, shutil, textwrap, requests, uuid
from bs4 import BeautifulSoup
import google.generativeai as genai
from sentence_transformers import SentenceTransformer
import chromadb
import gradio as gr
from langchain_community.document_loaders import UnstructuredPDFLoader
import camelot

# ======================
# 🔹 Gemini + Local Setup
# ======================
genai.configure(api_key="AQ.Ab8RN6I4za-ifPTOLhC78wesRWQxGaWeQQZZa44Cv7dBtaxX6A")  # 🔑 apni Gemini key
MODEL = "gemini-3.5-flash"

embedder = SentenceTransformer("all-MiniLM-L6-v2")
chroma_client = chromadb.Client()

# ======================
# 🔹 Utils
# ======================
def clean_text(text):
    return re.sub(r"\s+", " ", text).strip()

def adaptive_chunk_text(text):
    length = len(text)
    if length < 3000:
        size = 500
    elif length < 10000:
        size = 1000
    else:
        size = 1500
    chunks = []
    for i in range(0, len(text), size - 150):
        chunks.append(text[i:i + size])
    return chunks

def extract_pdf_text(pdf_path):
    """Smart PDF extractor (tables + text)"""
    full_text = ""
    try:
        tables = camelot.read_pdf(pdf_path, pages="all")
        for i, table in enumerate(tables):
            full_text += f"\n\n[Table {i+1}]\n" + table.df.to_string(index=False)
    except Exception:
        pass

    try:
        loader = UnstructuredPDFLoader(pdf_path)
        docs = loader.load()
        full_text += "\n\n".join([doc.page_content for doc in docs])
    except Exception as e:
        full_text += f"\n\n[Error extracting text: {e}]"

    return clean_text(full_text)

# ======================
# 🔹 Session Handling
# ======================
def create_user_collection():
    """Each user/session gets unique collection"""
    session_id = f"user_{str(uuid.uuid4())[:8]}"
    collection = chroma_client.create_collection(name=session_id)
    return session_id, collection

def reset_collection(collection_name):
    """Delete previous data for same user"""
    try:
        chroma_client.delete_collection(name=collection_name)
    except Exception:
        pass
    return chroma_client.create_collection(name=collection_name)

# ======================
# 🔹 Ingestion Logic
# ======================
def ingest_source(source, from_url, collection_name):
    # Delete previous user data
    collection = reset_collection(collection_name)

    if from_url:
        html = requests.get(source, timeout=15).text
        soup = BeautifulSoup(html, "html.parser")
        text = clean_text(soup.get_text())
    else:
        text = extract_pdf_text(source)

    if not text.strip():
        return "⚠️ No readable text found (maybe image-only PDF)."

    chunks = adaptive_chunk_text(text)
    embeddings = embedder.encode(chunks).tolist()

    for i, emb in enumerate(embeddings):
        collection.add(ids=[f"{collection_name}_{i}"], embeddings=[emb], documents=[chunks[i]])

    return f"✅ [{collection_name}] Ingested {len(chunks)} chunks successfully!"

# ======================
# 🔹 Query Logic
# ======================
def rag_query(query, collection_name):
    try:
        collection = chroma_client.get_collection(name=collection_name)
        q_emb = embedder.encode([query]).tolist()
        results = collection.query(query_embeddings=q_emb, n_results=4)
        if not results["documents"]:
            return "⚠️ No context found. Try ingesting data first."

        context = "\n\n".join(results["documents"][0])
        prompt = f"""
You are a knowledgeable AI assistant.
Use the context below to answer clearly and in multiple lines.

Context:
{context}

Question: {query}
Answer:
"""
        response = genai.GenerativeModel(MODEL).generate_content(prompt)
        ans = response.text.replace(". ", ".\n")
        return ans
    except Exception as e:
        return f"⚠️ Error: {e}"

# ======================
# 🔹 Gradio UI
# ======================
def start_new_session():
    session_id, _ = create_user_collection()
    return session_id

session_id = start_new_session()

def ingest_website(url):
    return ingest_source(url, True, session_id)

def ingest_pdf(file):
    return ingest_source(file.name, False, session_id)

def query_ask(q):
    return rag_query(q, session_id)

with gr.Blocks(theme=gr.themes.Soft(primary_hue="emerald")) as demo:
    gr.Markdown("# 🤖 Smart RAG App Pro (Gemini + Adaptive PDF + Multi-User Mode)")

    gr.Markdown(f"🆕 **Private Session ID:** `{session_id}` – Your data is isolated and auto-clears on refresh.")

    with gr.Tab("🌐 Ingest Website"):
        url_in = gr.Textbox(label="Enter Website URL")
        url_btn = gr.Button("Ingest Website")
        url_out = gr.Textbox(label="Status")
        url_btn.click(fn=ingest_website, inputs=url_in, outputs=url_out)

    with gr.Tab("📄 Ingest PDF"):
        pdf_in = gr.File(label="Upload PDF")
        pdf_btn = gr.Button("Ingest PDF")
        pdf_out = gr.Textbox(label="Status")
        pdf_btn.click(fn=ingest_pdf, inputs=pdf_in, outputs=pdf_out)

    with gr.Tab("💬 Ask Questions"):
        q_in = gr.Textbox(label="Ask anything from ingested data")
        q_btn = gr.Button("Ask Gemini")
        q_out = gr.Markdown(label="Answer")
        q_btn.click(fn=query_ask, inputs=q_in, outputs=q_out)

demo.launch()