from fastapi import APIRouter
from pydantic import BaseModel
import joblib
from sklearn.datasets import load_iris
from typing import Any, cast

from .config_huggingface import build_model_url, download_artifact_if_needed

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

class IrisFeatures(BaseModel):
    sepal_length: float = 1.0  
    sepal_width: float = 1.0
    petal_length: float = 1.0
    petal_width: float = 1.0


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

MODEL_URL = build_model_url("ML_DecisionTree_IrisClassifier.joblib")

iris = cast(Any, load_iris())


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.put("/models/irisClassifier")
def iris_classifier(data: IrisFeatures):
    import numpy as np

    try:
        _ensure_model_loaded()
    except Exception:
        detail = "Model not loaded."
        if MODEL_STATE["error"]:
            detail = f"Model not loaded: {MODEL_STATE['error']}"
        return {"error": detail, "status": 500}

    model = cast(Any, MODEL_STATE["model"])
    test_data = np.array([[
        data.sepal_length,
        data.sepal_width,
        data.petal_length,
        data.petal_width
    ]])
    raw_prediction = model.predict(test_data)[0]

    # Some serialized models output numeric indices, others output class labels.
    if isinstance(raw_prediction, (int, np.integer)):
        class_idx = int(raw_prediction)
        if class_idx < 0 or class_idx >= len(iris.target_names):
            return {"error": f"Invalid prediction index: {class_idx}", "status": 500}
        return {"prediction": str(iris.target_names[class_idx])}

    return {"prediction": str(raw_prediction)}