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 linear regression
class LinearRegressionRequest(BaseModel):
    age: int = 30
    monthly_spend: int = 50
    tenure_months: int = 12

from typing import Optional, Any

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

MODEL_URL = build_model_url("ML_LinearRegression_ChurnPredictor.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/linear_regression', summary="Predict churn risk with Bayesian Ridge Regression")
def predict_linear_regression(data: LinearRegressionRequest):
    import traceback
    try:
        _ensure_model_loaded()
    except Exception:
        detail = "Model not loaded."
        if MODEL_STATE["error"]:
            detail = f"Model not loaded: {MODEL_STATE['error']}"
        # Log traceback for debugging
        print("Model load error:", MODEL_STATE["error"])
        print(traceback.format_exc())
        return {"error": detail, "traceback": traceback.format_exc(), "status": 500}

    model = MODEL_STATE["model"]
    if model is None:
        # Defensive: model still not loaded
        error_msg = f"Model is None after loading. Error: {MODEL_STATE['error']}"
        print(error_msg)
        return {"error": error_msg, "status": 500}

    # Convert input data to DataFrame matching training features
    new_customer_data = pd.DataFrame([[data.age, data.monthly_spend, data.tenure_months]],
                                     columns=['age', 'monthly_spend', 'tenure_months'])
    # Make prediction
    try:
        pred_mean = model.predict(new_customer_data)[0]
    except Exception as e:
        print("Prediction error:", str(e))
        print(traceback.format_exc())
        return {"error": f"Prediction failed: {str(e)}", "traceback": traceback.format_exc(), "status": 500}

    percent = round(float(pred_mean) * 100, 1)
    if percent > 100:
        percent_str = '>100%'
    elif percent < 1:
        percent_str = '<1%'
    else:
        percent_str = f"{percent}%"
    return {
        "predicted_churn_risk_mean": percent_str
    }