from fastapi import APIRouter
from pydantic import BaseModel
import joblib
import pandas as pd

from .config_huggingface import build_model_url, download_artifact_if_needed

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

# Define the request model for logistic regression
class LogisticRegressionRequest(BaseModel):
    age: int = 30
    monthly_spend: float = 50.0
    tenure_months: int = 12

MODEL_STATE = {
    "model": None,
    "error": None,
}

MODEL_URL = build_model_url("ML_LogisticRegression_ChurnPredictor.joblib")

FEATURE_COLUMNS = ['age', 'monthly_spend', 'tenure_months']


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/logistic_regression', summary="Predict churn with Logistic Regression")
def predict_logistic_regression(data: LogisticRegressionRequest):
    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 = MODEL_STATE["model"]
    # Convert input data to DataFrame matching training features
    input_data = pd.DataFrame([[data.age, data.monthly_spend, data.tenure_months]], columns=FEATURE_COLUMNS)
    prediction = model.predict(input_data)[0]
    probability = model.predict_proba(input_data)[0].tolist()
    return {
        "prediction": int(prediction)
    }