"""
app.py - SQL Books RAG application powered by Gradio.
Retrieves relevant chunks from a FAISS index and generates answers
using Llama 3 via the Groq API (free tier).
"""

import json
import os

import faiss
import gradio as gr
import numpy as np
import requests
from sentence_transformers import SentenceTransformer

# -- Configuration -------------------------------------------------------------
INDEX_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "faiss_index")
EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
GEN_MODEL = "llama-3.1-8b-instant"
API_URL = "https://api.groq.com/openai/v1/chat/completions"
GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")
TOP_K = 5

# -- Load resources once at startup --------------------------------------------
print("Loading embedding model ...")
embedder = SentenceTransformer(EMBED_MODEL)

print("Loading FAISS index ...")
index = faiss.read_index(os.path.join(INDEX_DIR, "index.faiss"))

print("Loading chunk metadata ...")
with open(os.path.join(INDEX_DIR, "chunks.json"), "r", encoding="utf-8") as f:
    chunks = json.load(f)

print(f"Ready: {index.ntotal} vectors, {len(chunks)} chunks")
print(f"LLM: {GEN_MODEL} via Groq API")
print("App ready.")


# -- RAG pipeline ---------------------------------------------------------------
def retrieve(query: str, top_k: int = TOP_K):
    """Embed the query and retrieve the top-k most similar chunks."""
    query_vec = embedder.encode([query]).astype("float32")
    distances, indices = index.search(query_vec, top_k)
    results = []
    for dist, idx in zip(distances[0], indices[0]):
        if idx < len(chunks):
            results.append({
                "text": chunks[idx]["text"],
                "source": chunks[idx]["source"],
                "distance": float(dist),
            })
    return results


def generate_answer(query: str, context_chunks: list) -> str:
    """Build a prompt from retrieved context and generate an answer via Groq API."""
    context = "\n\n".join(
        f"[Source: {c['source']}]\n{c['text'][:800]}" for c in context_chunks
    )

    system_message = (
        "You are a helpful SQL tutor. Answer the user's question using ONLY the "
        "provided context from SQL textbooks. Give clear, detailed explanations with "
        "examples where appropriate. If the context doesn't contain enough information, "
        "say so honestly. Format your answer using markdown."
    )

    user_message = (
        f"## Context from SQL Textbooks\n\n{context}\n\n"
        f"---\n\n## Question\n{query}"
    )

    headers = {
        "Authorization": f"Bearer {GROQ_API_KEY}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": GEN_MODEL,
        "messages": [
            {"role": "system", "content": system_message},
            {"role": "user", "content": user_message},
        ],
        "max_tokens": 1024,
        "temperature": 0.3,
    }

    try:
        response = requests.post(API_URL, headers=headers, json=payload, timeout=60)
        response.raise_for_status()
        result = response.json()
        return result["choices"][0]["message"]["content"].strip()
    except requests.exceptions.HTTPError:
        return f"⚠️ API error ({response.status_code}): {response.text}"
    except Exception as e:
        return f"⚠️ Generation error: {e}"


def rag_query(question: str):
    """Full RAG pipeline: retrieve, generate, format output."""
    if not question.strip():
        return "Please enter a question.", ""

    # Retrieve
    retrieved = retrieve(question)

    # Generate
    answer = generate_answer(question, retrieved)

    # Format sources
    sources_text = "\n\n---\n\n".join(
        f"**Source {i+1}** - *{r['source']}*\n\n{r['text'][:500]}{'...' if len(r['text']) > 500 else ''}"
        for i, r in enumerate(retrieved)
    )

    return answer, sources_text


# -- Gradio UI ------------------------------------------------------------------
DESCRIPTION = """
# 📚 SQL Books RAG

Ask any question about SQL and get answers grounded in content from **5 SQL textbooks**:

- *Practical SQL: A Beginner's Guide to Storytelling with Data*
- *SQL for Data Scientists* (Renee M. Teate)
- *SQL for Data Analysis: Advanced Techniques for Transforming Data into Insights*
- *The Art of SQL*
- *Learning SQL: Generate, Manipulate, and Retrieve Data*

Powered by **FAISS** retrieval + **Llama 3.1** generation.
"""

EXAMPLES = [
    "What is a JOIN in SQL?",
    "Explain the difference between INNER JOIN and LEFT JOIN",
    "How do window functions work in SQL?",
    "What is a subquery and when should I use one?",
    "How do I use GROUP BY with HAVING?",
    "What are common table expressions (CTEs)?",
]

my_theme = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="blue",
)

with gr.Blocks(title="SQL Books RAG", theme=my_theme) as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column(scale=3):
            question = gr.Textbox(
                label="Your SQL Question",
                placeholder="e.g. What is a JOIN in SQL?",
                lines=2,
            )
            submit_btn = gr.Button("Ask", variant="primary", size="lg")
        with gr.Column(scale=1):
            gr.Markdown("### Try these examples")
            for ex in EXAMPLES:
                gr.Button(ex, size="sm").click(
                    fn=lambda e=ex: e, outputs=question
                )

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

    with gr.Accordion("Retrieved Source Chunks", open=False):
        sources_box = gr.Markdown()

    submit_btn.click(fn=rag_query, inputs=question, outputs=[answer_box, sources_box])
    question.submit(fn=rag_query, inputs=question, outputs=[answer_box, sources_box])


if __name__ == "__main__":
    demo.launch()