from fastapi import FastAPI, HTTPException, Depends, status
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from pydantic import BaseModel, Field
import joblib
import os
import numpy as np

# Initialize FastAPI app
app = FastAPI(
    title="Iris Classification API",
    description="A REST API for predicting Iris species using a pre-trained scikit-learn model.",
    version="1.0.0"
)

# --- Authentication Setup ---
security = HTTPBasic()

def get_current_username(credentials: HTTPBasicCredentials = Depends(security)):
    correct_username = os.getenv("API_USERNAME")
    correct_password = os.getenv("API_PASSWORD")

    if not correct_username or not correct_password:
        # This handles cases where secrets aren't set in HF Spaces (shouldn't happen if done correctly)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="API credentials not configured on the server."
        )

    if not (credentials.username == correct_username and credentials.password == correct_password):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect username or password",
            headers={"WWW-Authenticate": "Basic"},
        )
    return credentials.username

# --- Model Loading ---
model = None
class_names = None

@app.on_event("startup")
async def load_artifacts():
    global model, class_names
    model_path = os.path.join("model", "iris_model.joblib")
    class_names_path = os.path.join("model", "iris_class_names.joblib")

    if not os.path.exists(model_path) or not os.path.exists(class_names_path):
        raise RuntimeError(f"Model or class names file not found. Ensure '{model_path}' and '{class_names_path}' exist.")

    model = joblib.load(model_path)
    class_names = joblib.load(class_names_path)
    print("Model and class names loaded successfully.")

# --- Request Body Model ---
class IrisFeatures(BaseModel):
    sepal_length: float = Field(..., example=5.1, description="Sepal length in cm")
    sepal_width: float = Field(..., example=3.5, description="Sepal width in cm")
    petal_length: float = Field(..., example=1.4, description="Petal length in cm")
    petal_width: float = Field(..., example=0.2, description="Petal width in cm")

# --- API Endpoint ---
@app.post("/predict", summary="Predict Iris Species", response_description="The predicted Iris species and probabilities.")
async def predict_iris(
    features: IrisFeatures,
    current_user: str = Depends(get_current_username)
):
    if model is None or class_names is None:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Model is not loaded yet. Please try again in a moment."
        )

    input_data = np.array([[
        features.sepal_length,
        features.sepal_width,
        features.petal_length,
        features.petal_width
    ]])

    prediction_index = model.predict(input_data)[0]
    predicted_species = class_names[prediction_index]

    probabilities = model.predict_proba(input_data)[0]
    probabilities_dict = {name: float(prob) for name, prob in zip(class_names, probabilities)}

    return {
        "predicted_species": predicted_species,
        "prediction_probabilities": probabilities_dict
    }

# --- Health Check Endpoint ---
@app.get("/health", summary="Health Check", response_description="Indicates if the API is running.")
async def health_check():
    return {"status": "ok", "model_loaded": model is not None}

# Note: The uvicorn.run part is for local execution.
# Hugging Face Spaces will use the CMD in the Dockerfile.
# For local testing in Colab, you'd use ngrok or colabcode (see below).