# app.py
import gradio as gr

import spaces

from rag_agent import prepare_index_and_chunks, load_model
from sentence_transformers import SentenceTransformer
from utils.retrieval import retrieve_relevant_chunks
from utils.generation import generate_answer

# ——— FIXED CONFIGURATION ———
PDF_FOLDER      = "./data"  # your folder with all PDFs
EMBEDDER        = "all-MiniLM-L6-v2"
CHUNK_SIZE      = 500
OVERLAP         = 100
INDEX_TYPE      = "innerproduct"
MODEL_NAME      = "AliMaatouk/LLama-3-8B-Tele-it" #"deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
TOP_K           = 5

# ——— PREPARE INDEX & MODEL ONCE ———
faiss_index_path, chunks_path = prepare_index_and_chunks(
    pdf_folder=PDF_FOLDER,
    chunk_size=CHUNK_SIZE,
    overlap=OVERLAP,
    index_type=INDEX_TYPE,
    embedder_name=EMBEDDER
)
model, tokenizer = load_model(MODEL_NAME)
embedder = SentenceTransformer(EMBEDDER)

# ——— INFERENCE FUNCTION ———
@spaces.GPU()
def answer_query(query: str) -> str:
    if not query.strip():
        return "⚠️ Please enter a question."
    # Retrieve top-K chunks
    chunks = retrieve_relevant_chunks(
        query,
        embedder,
        TOP_K,
        faiss_index_path,
        chunks_path
    )
    # Generate answer
    return generate_answer(query, chunks, model, tokenizer)

# ——— GRADIO UI ———
iface = gr.Interface(
    fn=answer_query,
    inputs=gr.Textbox(lines=2, placeholder="Type your question here…", label="Question"),
    outputs=gr.Textbox(label="Answer"),
    title="📡 SpectrumGPT",
    description=(
        "Answer questions on spectrum regulations.\n\n"
    )
)

if __name__ == "__main__":
    iface.launch(share=True)