# -*- coding: utf-8 -*-
"""Untitled7.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1nmczTgIa8AsM0b0qWnhNbUavzSKsMm_m
"""

!pip -q install gradio groq sentence-transformers faiss-cpu pypdf python-docx pandas

import os, re
import gradio as gr
import pandas as pd

from groq import Groq
from sentence_transformers import SentenceTransformer
import faiss

from pypdf import PdfReader
from docx import Document
from google.colab import userdata

# ----------------------------
# Load Groq key from Colab Secrets
# ----------------------------
os.environ["GROQ_API_KEY"] = userdata.get("GROQ_API_KEY")
assert os.environ["GROQ_API_KEY"], "❌ GROQ_API_KEY missing in Colab Secrets"

# ----------------------------
# Config
# ----------------------------
DEFAULT_MODEL = "llama-3.1-8b-instant"
EMB_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
CHUNK_SIZE = 900
CHUNK_OVERLAP = 150

# ----------------------------
# File Readers
# ----------------------------
def read_pdf(path):
    return "\n".join([p.extract_text() or "" for p in PdfReader(path).pages])

def read_docx(path):
    return "\n".join(p.text for p in Document(path).paragraphs)

def read_txt(path):
    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        return f.read()

def read_csv(path):
    return pd.read_csv(path).to_csv(index=False)

def load_file(path):
    ext = os.path.splitext(path)[1].lower()
    if ext == ".pdf": return read_pdf(path)
    if ext == ".docx": return read_docx(path)
    if ext == ".csv": return read_csv(path)
    return read_txt(path)

# ----------------------------
# Chunking
# ----------------------------
def chunk_text(text):
    text = re.sub(r"\s+", " ", text)
    chunks, start = [], 0
    while start < len(text):
        end = start + CHUNK_SIZE
        chunks.append(text[start:end])
        start = end - CHUNK_OVERLAP
    return chunks

# ----------------------------
# Build FAISS index
# ----------------------------
def build_index(files):
    emb = SentenceTransformer(EMB_MODEL)
    texts, meta = [], []

    for f in files:
        raw = load_file(f.name)
        for i, chunk in enumerate(chunk_text(raw)):
            texts.append(chunk)
            meta.append({"file": os.path.basename(f.name), "chunk": i})

    vectors = emb.encode(texts, convert_to_numpy=True)
    faiss.normalize_L2(vectors)

    index = faiss.IndexFlatIP(vectors.shape[1])
    index.add(vectors)

    return {
        "index": index,
        "emb": emb,
        "texts": texts,
        "meta": meta
    }, f"✅ Indexed {len(texts)} chunks"

# ----------------------------
# Retrieval
# ----------------------------
def retrieve(store, query, k=5):
    qv = store["emb"].encode([query], convert_to_numpy=True)
    faiss.normalize_L2(qv)
    scores, ids = store["index"].search(qv, k)

    return [
        {
            "text": store["texts"][i],
            "meta": store["meta"][i],
            "score": float(s)
        }
        for i, s in zip(ids[0], scores[0]) if i != -1
    ]

# ----------------------------
# Groq LLM
# ----------------------------
def ask_groq(question, contexts):
    client = Groq()

    context_text = "\n\n".join(
        f"[{c['meta']['file']} | chunk {c['meta']['chunk']}]\n{c['text']}"
        for c in contexts
    )

    messages = [
        {"role": "system", "content": "Answer strictly from context. Cite sources."},
        {"role": "user", "content": f"Context:\n{context_text}\n\nQuestion:\n{question}"}
    ]

    response = client.chat.completions.create(
        model=DEFAULT_MODEL,
        messages=messages,
        temperature=0.2
    )
    return response.choices[0].message.content

# ----------------------------
# Gradio Handlers
# ----------------------------
def index_files(files):
    store, msg = build_index(files)
    return store, msg

def chat(q, store):
    if store is None:
        return "❌ Upload and index files first"
    ctx = retrieve(store, q)
    return ask_groq(q, ctx)

# ----------------------------
# UI
# ----------------------------
with gr.Blocks() as app:
    gr.Markdown("# 📚 Groq RAG Application")

    store = gr.State(None)
    files = gr.File(file_count="multiple", label="Upload documents")

    build = gr.Button("Build Index")
    status = gr.Textbox(label="Status")

    q = gr.Textbox(label="Question")
    a = gr.Markdown()

    build.click(index_files, files, [store, status])
    gr.Button("Ask").click(chat, [q, store], a)

app.launch(share=True)