from fastapi import APIRouter
from pydantic import BaseModel
from backend.app.model import predict

router = APIRouter()

RISK_LABELS = {0: "Low Risk", 1: "Medium Risk", 2: "High Risk"}

class InputData(BaseModel):
    Age: int
    SystolicBP: int
    DiastolicBP: int
    BS: float
    BodyTemp: float
    HeartRate: int

@router.post("/predict")
def predict_risk(input_data: InputData):
    input_str = (
        f"Age: {input_data.Age}, "
        f"SystolicBP: {input_data.SystolicBP}, "
        f"DiastolicBP: {input_data.DiastolicBP}, "
        f"BS: {input_data.BS}, "
        f"BodyTemp: {input_data.BodyTemp}, "
        f"HeartRate: {input_data.HeartRate}. "
        f"Predict the Risk Level."
    )
    prediction_index = predict(input_str)  # ✅ now returns index
    return {"risk_level": RISK_LABELS.get(prediction_index, "Unknown")}