from fastapi import APIRouter
from pydantic import BaseModel
import joblib
import pandas as pd
from typing import Optional, Any

from .config_huggingface import build_model_url, download_artifact_if_needed

router = APIRouter(tags=["Machine Learning"])


class MeanShiftRequest(BaseModel):
    session_duration_min: float = 18.0
    pages_viewed: int = 10
    items_purchased: int = 3


MODEL_STATE: dict[str, Optional[Any]] = {
    "model": None,
    "error": None,
}

MODEL_URL = build_model_url("ML_MeanShift_UserBehavior.joblib")


def _ensure_model_loaded() -> None:
    if MODEL_STATE["model"] is not None:
        return
    try:
        model_path = download_artifact_if_needed(MODEL_URL)
        MODEL_STATE["model"] = joblib.load(model_path)
        MODEL_STATE["error"] = None
    except Exception as e:
        MODEL_STATE["error"] = str(e)
        raise


@router.post("/models/mean_shift", summary="Classify user behaviour with Mean Shift clustering")
def predict_mean_shift(data: MeanShiftRequest):
    import traceback
    try:
        _ensure_model_loaded()
    except Exception:
        detail = f"Model not loaded: {MODEL_STATE['error']}" if MODEL_STATE["error"] else "Model not loaded."
        return {"error": detail, "traceback": traceback.format_exc(), "status": 500}

    bundle = MODEL_STATE["model"]
    if bundle is None:
        return {"error": f"Model is None. Error: {MODEL_STATE['error']}", "status": 500}

    scaler           = bundle["scaler"]
    ms               = bundle["ms"]
    cluster_name_map = bundle["cluster_name_map"]

    input_df = pd.DataFrame(
        [[data.session_duration_min, data.pages_viewed, data.items_purchased]],
        columns=["session_duration_min", "pages_viewed", "items_purchased"],
    )

    try:
        cluster = int(ms.predict(scaler.transform(input_df))[0])
    except Exception as e:
        return {"error": f"Prediction failed: {str(e)}", "traceback": traceback.format_exc(), "status": 500}

    return {
        "cluster": cluster,
        "label": cluster_name_map.get(cluster, f"Cluster {cluster}"),
    }