from fastapi import FastAPI
from pydantic import BaseModel

import xgboost as xgb
import pandas as pd

from huggingface_hub import hf_hub_download

app = FastAPI()

# =========================
# DOWNLOAD MODEL
# =========================

MODEL_PATH = hf_hub_download(
    repo_id="mjpsm/communicative-engagement-model",
    filename="engagement_xgb_model.json"
)

# =========================
# LOAD MODEL
# =========================

model = xgb.XGBClassifier()

model.load_model(MODEL_PATH)

# =========================
# FEATURES
# =========================

FEATURES = [
    "total_time",
    "was_webcam_on",
    "screenshare_usage",
    "never_spoke",
    "speech_turns"
]

# =========================
# LABEL MAP
# =========================

LABEL_MAP = {
    0: "Silent Observer",
    1: "Occasional Participant",
    2: "Active Participant"
}

# =========================
# REQUEST MODEL
# =========================

class PredictionRequest(BaseModel):
    rows: list

# =========================
# HEALTH CHECK
# =========================

@app.get("/")
def home():

    return {
        "status": "running",
        "model": "mjpsm/communicative-engagement-model"
    }

# =========================
# PREDICTION ENDPOINT
# =========================

@app.post("/predict")
def predict(request: PredictionRequest):

    df = pd.DataFrame(request.rows)

    # Ensure all required features exist
    for col in FEATURES:

        if col not in df:
            df[col] = 0

    # =========================
    # FEATURE EXTRACTION
    # =========================

    X = df[FEATURES]

    # =========================
    # MODEL PREDICTION
    # =========================

    preds = model.predict(X)

    # =========================
    # LABEL MAPPING
    # =========================

    df["cluster_label"] = [
        LABEL_MAP[p]
        for p in preds
    ]

    # =========================
    # FINAL OUTPUT
    # =========================

    output = df[[
        "student_name",
        "meeting_id",
        "cluster_label"
    ]]

    return output.to_dict(
        orient="records"
    )