import os
import random
import torch
import pandas as pd
from functools import lru_cache

from flask import Flask, render_template, request, jsonify
from sentence_transformers import SentenceTransformer

from fastapi import FastAPI
from fastapi.middleware.wsgi import WSGIMiddleware
import uvicorn

# ===============================
# CONFIG HF
# ===============================
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
CSV_DATA = os.path.join(BASE_DIR, "dataset_2026.csv")
EMB_FILE = os.path.join(BASE_DIR, "embeddings_questions.pt")

TOP_K_RECOMMANDATIONS = 5
DEVICE = "cpu"   # ⛔ FORCÉ CPU (HF)

# ===============================
# FLASK APP
# ===============================
app = Flask(
    __name__,
    template_folder=os.path.join(BASE_DIR, "templates"),
    static_folder=os.path.join(BASE_DIR, "static")
)

# ===============================
# LOAD MODEL (SAFE)
# ===============================
print("🔹 Loading model (CPU only)...")
model = SentenceTransformer(
    "sentence-transformers/all-MiniLM-L6-v2",
    device="cpu"
)

# ===============================
# LOAD DATASET
# ===============================
print("🔹 Loading dataset...")
df = pd.read_csv(CSV_DATA)
df = df.dropna(subset=["question"]).reset_index(drop=True)

questions = df["question"].astype(str).tolist()
NB_QUESTIONS = len(questions)

print(f"✅ Questions loaded: {NB_QUESTIONS}")

# ===============================
# LOAD / CREATE EMBEDDINGS
# ===============================
if os.path.exists(EMB_FILE):
    print("🔹 Loading cached embeddings...")
    emb_base = torch.load(EMB_FILE, map_location="cpu")
else:
    print("🔹 Computing embeddings...")
    emb_base = model.encode(
        questions,
        convert_to_tensor=True,
        normalize_embeddings=True,
        batch_size=32
    )
    torch.save(emb_base, EMB_FILE)

emb_base = emb_base.cpu()

# ===============================
# CACHE QUESTION EMBEDDING
# ===============================
@lru_cache(maxsize=500)
def encode_question_cached(q: str):
    return model.encode(
        q,
        convert_to_tensor=True,
        normalize_embeddings=True
    ).cpu()

# ===============================
# UTILS
# ===============================
def enrich_message(base):
    return random.choice([
        f"Bonne question 🙂 {base}",
        f"Voici ce que je peux vous dire : {base}",
        base
    ])

# ===============================
# CORE LOGIC
# ===============================
def process_question(question: str):

    if not question.strip():
        return {"response": "Veuillez poser une question."}

    emb_q = encode_question_cached(question).unsqueeze(0)

    scores = torch.matmul(emb_q, emb_base.T).squeeze(0)
    values, indices = torch.topk(scores, k=min(TOP_K_RECOMMANDATIONS + 1, len(scores)))

    best_idx = indices[0].item()
    confidence = int(values[0].item() * 100)

    if confidence < 50:
        return {
            "response": "Je ne suis pas sûr de la réponse.",
            "confidence": confidence
        }

    return {
        "response": enrich_message(df["rationale"].iloc[best_idx]),
        "confidence": confidence,
        "matched": df["question"].iloc[best_idx],
        "intent": df["intent"].iloc[best_idx]
    }

# ===============================
# ROUTES
# ===============================
@app.route("/")
def index():
    return render_template("index.html")

@app.route("/ask", methods=["POST"])
def ask():
    data = request.get_json() or {}
    question = data.get("question", "")
    return jsonify(process_question(question))

# ===============================
# FASTAPI WRAPPER (HF)
# ===============================
fastapi_app = FastAPI()
fastapi_app.mount("/", WSGIMiddleware(app))

# ===============================
# MAIN
# ===============================
if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    print(f"🚀 Running on port {port}")
    uvicorn.run(fastapi_app, host="0.0.0.0", port=port)