import os
from fastapi import FastAPI
from pydantic import BaseModel
from utils.ml import train_model, save_model, load_model, predict

app = FastAPI()
model = None

# Request body for prediction
class IrisFeatures(BaseModel):
    sepal_length: float
    sepal_width: float
    petal_length: float
    petal_width: float

# Ensure the model is loaded on startup
@app.on_event("startup")
def startup_event():
    if not os.path.exists("ml_models/iris.model"):
        _model = train_model()
        did_save_model = save_model(model=_model)
        if did_save_model:
            print("Model trained and saved successfully.")
        else:
            print("Model training and saving failed.")
    global model
    model = load_model()

# The prediction endpoint
@app.post("/predict")
def make_prediction(iris: IrisFeatures):
    return predict(
        model=model,
        sepal_length=iris.sepal_length,
        sepal_width=iris.sepal_width,
        petal_length=iris.petal_length,
        petal_width=iris.petal_width
    )