from __future__ import annotations

from pathlib import Path
from typing import Any, Union

import joblib
import pandas as pd
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

MODEL_PATH = Path(__file__).resolve().parent.parent / "xg_bost.pkl"

# LabelEncoder-style mappings (alphabetical order) for categorical features.
CATEGORICAL_MAPPINGS = {
    "weather_condition": {
        "Clear": 0.0,
        "Cloudy": 1.0,
        "Fog": 2.0,
        "Rain": 3.0,
        "Snow": 4.0,
        "Windy": 5.0,
    },
    "traffic_level": {
        "High": 0.0,
        "Low": 1.0,
        "Medium": 2.0,
        "Very High": 3.0,
        "Very Low": 4.0,
    },
    "vehicle_type": {
        "car": 0.0,
        "cycle": 1.0,
        "ev": 2.0,
        "motorcycle": 3.0,
        "scooter": 4.0,
    },
    "order_type": {
        "Buffet": 0.0,
        "Drinks": 1.0,
        "Meal": 2.0,
        "Snack": 3.0,
        "Unknown": 4.0,
    },
}

model = joblib.load(MODEL_PATH)
FEATURE_ORDER = list(model.feature_names_in_)


class PredictInput(BaseModel):
    distance_km: float = Field(..., example=7.93)
    weather_condition: Union[str, float, int] = Field(..., example="Windy")
    traffic_level: Union[str, float, int] = Field(..., example="Low")
    vehicle_type: Union[str, float, int] = Field(..., example="scooter")
    temperature_c: float = Field(..., example=23.0)
    humidity_pct: float = Field(..., example=55.0)
    precipitation_mm: float = Field(..., example=0.0)
    preparation_time_min: float = Field(..., example=12.0)
    courier_experience_yrs: float = Field(..., example=1.0)
    worker_age: float = Field(..., example=29.0)
    worker_rating: float = Field(..., example=4.7)
    order_type: Union[str, float, int] = Field(..., example="Unknown")
    weather_risk: float = Field(..., example=7.0)
    traffic_risk: float = Field(..., example=25.0)


def _encode_categorical(feature_name: str, value: Any) -> float:
    if isinstance(value, (int, float)) and not isinstance(value, bool):
        return float(value)

    mapping = CATEGORICAL_MAPPINGS[feature_name]
    raw = str(value).strip()
    if raw in mapping:
        return mapping[raw]

    lower_map = {k.lower(): v for k, v in mapping.items()}
    if raw.lower() in lower_map:
        return lower_map[raw.lower()]

    allowed = ", ".join(mapping.keys())
    raise HTTPException(
        status_code=422,
        detail=f"Unknown value '{value}' for '{feature_name}'. Allowed values: {allowed}",
    )


def _prepare_features(payload: PredictInput) -> pd.DataFrame:
    values = payload.model_dump()
    row: dict[str, float] = {}

    for feature in FEATURE_ORDER:
        value = values[feature]
        if feature in CATEGORICAL_MAPPINGS:
            row[feature] = _encode_categorical(feature, value)
        else:
            row[feature] = float(value)

    return pd.DataFrame([row], columns=FEATURE_ORDER)


app = FastAPI(title="Gitwire XGBoost Inference API", version="1.0.0")


@app.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@app.post("/predit")
def predit(payload: PredictInput) -> dict[str, float]:
    features = _prepare_features(payload)
    prediction = float(model.predict(features)[0])
    return {"prediction": prediction}