from flask import Flask, request, jsonify, render_template
import os, json, numpy as np
from pathlib import Path
import faiss
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from sentence_transformers import SentenceTransformer
from .llm_provider import answer_with_llm

BASE = Path(__file__).resolve().parent.parent
INDEX = BASE / "rag" / "vector.index"
CHUNKS = BASE / "rag" / "chunks.jsonl"

app = Flask(
    __name__,
    template_folder=str(BASE / "web" / "templates"),
    static_folder=str(BASE / "web" / "static")
)

embedder = None
reranker_model = None
reranker_tok = None
index = None
texts = None

device = "cuda" if torch.cuda.is_available() else "cpu"


def expand_queries(question):
    try:
        prompt = f"""
Buat 8 variasi pertanyaan lain dari pertanyaan berikut:

"{question}"

Format output: list tanpa nomor.
"""
        res = answer_with_llm(prompt, "")
        qs = [question]
        for ln in res.split("\n"):
            if "-" in ln:
                qs.append(ln.replace("-", "").strip())
        return qs[:10]
    except:
        return [question]

def compress_context(chunks):
    ctx = "\n\n".join(chunks)
    if len(ctx) < 2500:
        return ctx
    try:
        return answer_with_llm(
            "Ringkas konteks berikut tanpa mengubah makna.",
            ctx
        )
    except:
        return ctx[:3000]


def ensure_loaded():
    global embedder, reranker_model, reranker_tok, index, texts

    if embedder is None:
        embedder = SentenceTransformer("intfloat/multilingual-e5-large-instruct")

    if reranker_model is None:
        reranker_tok = AutoTokenizer.from_pretrained("BAAI/bge-reranker-v2-m3")
        reranker_model = AutoModelForSequenceClassification.from_pretrained(
            "BAAI/bge-reranker-v2-m3"
        ).to(device)

    if index is None:
        if not INDEX.exists():
            raise FileNotFoundError("FAISS index tidak ditemukan. Jalankan rag/embed.py.")
        index = faiss.read_index(str(INDEX))

    if texts is None:
        with open(CHUNKS, "r", encoding="utf-8") as f:
            texts = [json.loads(l)["text"] for l in f if l.strip()]


@app.route("/")
def home():
    return render_template("index.html")


@app.route("/query", methods=["POST"])
def query_api():
    data = request.get_json(force=True)
    question = data.get("question", "")
    use_llm = bool(data.get("use_llm", True))

    top_k = int(data.get("top_k", 5))
    retrieve_k = max(top_k * 5, 25)

    if not question.strip():
        return jsonify({"error": "Pertanyaan kosong"}), 400

    ensure_loaded()

    # Multi-query expansion
    queries = expand_queries(question)

    # Gather candidates
    candidates = {}
    for q in queries:
        q_emb = embedder.encode([f"query: {q}"], convert_to_numpy=True).astype("float32")
        faiss.normalize_L2(q_emb)
        D, I = index.search(q_emb, retrieve_k)
        for idx in I[0]:
            candidates[idx] = texts[idx]

    # Reranking
    scored = []
    for idx, text in candidates.items():
        inp = reranker_tok(
            question, text,
            truncation=True, padding=True,
            max_length=512, return_tensors="pt"
        ).to(device)
        with torch.no_grad():
            s = reranker_model(**inp).logits.squeeze().item()
        scored.append((s, idx, text))

    ranked = sorted(scored, reverse=True)

    hits = [{"score": float(s), "text": text} for s, _, text in ranked[:top_k]]

    # Optional LLM
    answer = None
    if use_llm:
        ctx = compress_context([h["text"] for h in hits])
        answer = answer_with_llm(question, ctx)

    return jsonify({
        "question": question,
        "hits": hits,
        "answer": answer
    })


if __name__ == "__main__":
    app.run(debug=True)
