from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
from huggingface_hub import hf_hub_download
import joblib
import json
import re
import os

app = FastAPI(title="On-Task Participation Detector")

# ── Load model files from HuggingFace Hub ────────────────────
REPO_ID = "ethnmcl/ontask-participation-classifier"

print("Loading model files...")
clf_path      = hf_hub_download(repo_id=REPO_ID, filename="classifier.pkl")
metadata_path = hf_hub_download(repo_id=REPO_ID, filename="metadata.json")

clf      = joblib.load(clf_path)
metadata = json.load(open(metadata_path))

THRESHOLD = metadata.get("threshold", 0.40)

print("Loading embedder...")
embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

print(f"✅ Model loaded — threshold: {THRESHOLD}")

# ── Utils ─────────────────────────────────────────────────────
def clean_text(text: str) -> str:
    text = str(text).strip()
    text = re.sub(r'http\S+|www\S+', ' ', text)
    text = re.sub(r'<[^>]+>', ' ', text)
    text = re.sub(r'\s+', ' ', text)
    return text.strip()

# ── Request / Response schemas ────────────────────────────────
class CheckinRequest(BaseModel):
    yesterday_checkin: str
    today_checkin:     str

class CheckinResponse(BaseModel):
    on_task:        bool
    confidence:     float
    threshold_used: float
    label:          str

# ── Routes ────────────────────────────────────────────────────
@app.get("/")
def root():
    return {
        "name":        "On-Task Participation Detector",
        "repo":        REPO_ID,
        "threshold":   THRESHOLD,
        "auc_roc":     metadata.get("auc_roc"),
        "status":      "ready"
    }

@app.get("/health")
def health():
    return {"status": "ok"}

@app.post("/predict", response_model=CheckinResponse)
def predict(req: CheckinRequest):
    if not req.today_checkin.strip():
        raise HTTPException(status_code=400, detail="today_checkin cannot be empty")

    text  = f"Yesterday: {clean_text(req.yesterday_checkin)} Today: {clean_text(req.today_checkin)}"
    emb   = embedder.encode([text])
    proba = float(clf.predict_proba(emb)[0][1])
    label = proba >= THRESHOLD

    return CheckinResponse(
        on_task        = bool(label),
        confidence     = round(proba, 3),
        threshold_used = THRESHOLD,
        label          = "on_task" if label else "not_on_task"
    )

@app.post("/predict/batch")
def predict_batch(requests: list[CheckinRequest]):
    if len(requests) > 100:
        raise HTTPException(status_code=400, detail="Max 100 requests per batch")

    texts  = [
        f"Yesterday: {clean_text(r.yesterday_checkin)} Today: {clean_text(r.today_checkin)}"
        for r in requests
    ]
    embs   = embedder.encode(texts)
    probas = clf.predict_proba(embs)[:, 1]

    return [
        {
            "on_task":        bool(p >= THRESHOLD),
            "confidence":     round(float(p), 3),
            "threshold_used": THRESHOLD,
            "label":          "on_task" if p >= THRESHOLD else "not_on_task"
        }
        for p in probas
    ]