# =======================================
# 📘 RAG App Pro – Groq + Smart Embeddings
#    Multi-User Safe
# =======================================


import os
import re
import requests
import uuid

from bs4 import BeautifulSoup
from groq import Groq

from sentence_transformers import SentenceTransformer
import chromadb
import gradio as gr

from langchain_community.document_loaders import UnstructuredPDFLoader
import camelot


# ======================
# 🔹 Groq + Local Setup
# ======================

# Set your Groq API key as an environment variable:
# Windows CMD:
# set GROQ_API_KEY=your_groq_api_key
#
# PowerShell:
# $env:GROQ_API_KEY="your_groq_api_key"

GROQ_API_KEY = os.getenv("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")

if not GROQ_API_KEY:
    raise ValueError(
        "GROQ_API_KEY not found. Please set your Groq API key "
        "as an environment variable."
    )

client = Groq(api_key=GROQ_API_KEY)

# Groq model
MODEL = "llama-3.3-70b-versatile"

# Local embedding model
embedder = SentenceTransformer("all-MiniLM-L6-v2")

# ChromaDB
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 = ""

    # ----------------------
    # Extract tables
    # ----------------------
    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

    # ----------------------
    # Extract normal text
    # ----------------------
    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 a 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 user data"""

    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
    )

    # ----------------------
    # Website
    # ----------------------
    if from_url:

        try:
            response = requests.get(
                source,
                timeout=15,
                headers={
                    "User-Agent": "Mozilla/5.0"
                }
            )

            response.raise_for_status()

            soup = BeautifulSoup(
                response.text,
                "html.parser"
            )

            # Remove unnecessary HTML
            for tag in soup(
                ["script", "style", "noscript"]
            ):
                tag.decompose()

            text = clean_text(
                soup.get_text()
            )

        except Exception as e:

            return f"⚠️ Website error: {e}"

    # ----------------------
    # PDF
    # ----------------------
    else:

        try:
            text = extract_pdf_text(source)

        except Exception as e:

            return f"⚠️ PDF error: {e}"

    # ----------------------
    # Check text
    # ----------------------

    if not text.strip():

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

    # ----------------------
    # Chunking
    # ----------------------

    chunks = adaptive_chunk_text(text)

    # ----------------------
    # Embeddings
    # ----------------------

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

    # ----------------------
    # Store in ChromaDB
    # ----------------------

    for i, emb in enumerate(embeddings):

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

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


# ======================
# 🔹 Query Logic
# ======================

def rag_query(query, collection_name):

    try:

        # ----------------------
        # Get collection
        # ----------------------

        collection = chroma_client.get_collection(
            name=collection_name
        )

        # ----------------------
        # Embed user query
        # ----------------------

        q_emb = embedder.encode(
            [query]
        ).tolist()

        # ----------------------
        # Semantic search
        # ----------------------

        results = collection.query(
            query_embeddings=q_emb,
            n_results=4
        )

        if (
            not results["documents"]
            or not results["documents"][0]
        ):
            return (
                "⚠️ No context found. "
                "Try ingesting data first."
            )

        # ----------------------
        # Build context
        # ----------------------

        context = "\n\n".join(
            results["documents"][0]
        )

        # ----------------------
        # Groq prompt
        # ----------------------

        system_prompt = """
You are a knowledgeable RAG assistant.

Answer the user's question using ONLY
the provided context.

Rules:
1. Do not invent information.
2. If the answer is not present in the context,
   clearly say that the information was not found.
3. Give a clear and useful answer.
4. Use multiple lines when appropriate.
5. Keep the answer concise but informative.
"""

        user_prompt = f"""
Context:

{context}

Question:

{query}
"""

        # ----------------------
        # Groq API
        # ----------------------

        response = client.chat.completions.create(

            model=MODEL,

            messages=[
                {
                    "role": "system",
                    "content": system_prompt
                },
                {
                    "role": "user",
                    "content": user_prompt
                }
            ],

            temperature=0.2,

            max_completion_tokens=1000
        )

        # ----------------------
        # Get answer
        # ----------------------

        answer = response.choices[
            0
        ].message.content

        return answer

    except Exception as e:

        return f"⚠️ Error: {e}"


# ======================
# 🔹 Session
# ======================

def start_new_session():

    session_id, _ = create_user_collection()

    return session_id


session_id = start_new_session()


# ======================
# 🔹 Gradio Functions
# ======================

def ingest_website(url):

    return ingest_source(
        url,
        True,
        session_id
    )


def ingest_pdf(file):

    if file is None:
        return "⚠️ Please upload a PDF."

    return ingest_source(
        file.name,
        False,
        session_id
    )


def query_ask(q):

    if not q.strip():
        return "⚠️ Please enter a question."

    return rag_query(
        q,
        session_id
    )


# ======================
# 🔹 Gradio UI
# ======================

with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue="emerald"
    )
) as demo:

    gr.Markdown(
        "# 🤖 Smart RAG App Pro "
        "(Groq + Adaptive PDF + Multi-User Mode)"
    )

    gr.Markdown(
        f"""
🆕 **Private Session ID:** `{session_id}`

Your uploaded/ingested data is isolated
inside your session.
"""
    )

    # ======================
    # Website Tab
    # ======================

    with gr.Tab("🌐 Ingest Website"):

        url_in = gr.Textbox(
            label="Enter Website URL",
            placeholder="https://example.com"
        )

        url_btn = gr.Button(
            "Ingest Website"
        )

        url_out = gr.Textbox(
            label="Status"
        )

        url_btn.click(
            fn=ingest_website,
            inputs=url_in,
            outputs=url_out
        )

    # ======================
    # PDF Tab
    # ======================

    with gr.Tab("📄 Ingest PDF"):

        pdf_in = gr.File(
            label="Upload PDF",
            file_types=[".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
        )

    # ======================
    # Ask Tab
    # ======================

    with gr.Tab("💬 Ask Questions"):

        q_in = gr.Textbox(
            label="Ask anything from ingested data",
            placeholder="Ask a question..."
        )

        q_btn = gr.Button(
            "Ask Groq"
        )

        q_out = gr.Markdown(
            label="Answer"
        )

        q_btn.click(
            fn=query_ask,
            inputs=q_in,
            outputs=q_out
        )


# ======================
# 🔹 Launch
# ======================

demo.launch()