app.py

from __future__ import annotations

import os
import sys
import logging
from pathlib import Path
from typing import List

import joblib
import numpy as np
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='[%(asctime)s] %(name)s — %(levelname)s: %(message)s'
)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="CyHub Model 4 — Domain Classification",
    description="Multi-class web data classifier (Normal/Adult/Betting/Malware)",
    version="1.0.0"
)

# CORS configuration
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ─────────────────────────────────────────────────────────────────
# Configuration
# ─────────────────────────────────────────────────────────────────

# Expected feature count (must match training data)
EXPECTED_FEATURES = 5

# Model file paths (support both relative and absolute)
MODEL_DIR = Path(os.getenv("MODEL_DIR", "."))
MODEL_PATH = MODEL_DIR / "trained_lightgbm_model.pkl"
ENCODER_PATH = MODEL_DIR / "label_encoder.pkl"

logger.info(f"Looking for model at: {MODEL_PATH.absolute()}")
logger.info(f"Looking for encoder at: {ENCODER_PATH.absolute()}")

# ─────────────────────────────────────────────────────────────────
# Request/Response Models
# ─────────────────────────────────────────────────────────────────

class InputData(BaseModel):
    """Input data for model prediction.

    Must be preprocessed (scaled, imputed, encoded) to match training data.
    """
    features: List[float] = Field(
        ...,
        min_items=EXPECTED_FEATURES,
        max_items=EXPECTED_FEATURES,
        description=f"Exactly {EXPECTED_FEATURES} preprocessed float values"
    )


class PredictionResponse(BaseModel):
    """Model prediction response."""
    predicted_label: str
    raw_prediction_encoded: int


# ─────────────────────────────────────────────────────────────────
# Model Loading
# ─────────────────────────────────────────────────────────────────

model = None
label_encoder = None
model_loaded = False

def load_models():
    """Load LightGBM model and label encoder."""
    global model, label_encoder, model_loaded

    try:
        # Check if files exist
        if not MODEL_PATH.exists():
            raise FileNotFoundError(f"Model file not found: {MODEL_PATH.absolute()}")
        if not ENCODER_PATH.exists():
            raise FileNotFoundError(f"Encoder file not found: {ENCODER_PATH.absolute()}")

        # Load model
        logger.info(f"Loading model from {MODEL_PATH.absolute()}...")
        model = joblib.load(str(MODEL_PATH))
        logger.info("✓ Model loaded successfully")

        # Load label encoder
        logger.info(f"Loading label encoder from {ENCODER_PATH.absolute()}...")
        label_encoder = joblib.load(str(ENCODER_PATH))
        logger.info("✓ Label encoder loaded successfully")

        # Verify encoder has classes
        if not hasattr(label_encoder, 'classes_'):
            raise ValueError("Label encoder missing 'classes_' attribute")

        logger.info(f"Label classes: {label_encoder.classes_.tolist()}")
        model_loaded = True

        return True

    except FileNotFoundError as e:
        logger.error(f"File error: {e}")
        logger.warning(f"Make sure model files are in: {MODEL_DIR.absolute()}")
        return False
    except Exception as e:
        logger.error(f"Error loading models: {e}")
        import traceback
        traceback.print_exc()
        return False


# ─────────────────────────────────────────────────────────────────
# Startup/Shutdown Events
# ─────────────────────────────────────────────────────────────────

@app.on_event("startup")
async def startup_event():
    """Initialize models on startup."""
    global model_loaded
    logger.info("Starting up CyHub Model 4 API...")

    if not load_models():
        logger.error("Failed to load models. API will return 503 until models are available.")
        model_loaded = False
    else:
        logger.info("Model 4 API ready!")


@app.on_event("shutdown")
async def shutdown_event():
    """Cleanup on shutdown."""
    logger.info("Shutting down CyHub Model 4 API...")


# ─────────────────────────────────────────────────────────────────
# API Endpoints
# ─────────────────────────────────────────────────────────────────

@app.get("/")
async def root():
    """Root endpoint — health check."""
    return {
        "service": "CyHub Model 4 — Domain Classification",
        "status": "healthy",
        "model_loaded": model_loaded,
        "version": "1.0.0"
    }


@app.get("/health")
async def health_check():
    """Health check endpoint."""
    if not model_loaded:
        raise HTTPException(
            status_code=503,
            detail="Model not loaded. Check server logs."
        )

    return {
        "status": "healthy",
        "model_status": "ready",
        "expected_features": EXPECTED_FEATURES
    }


@app.post("/predict", response_model=PredictionResponse)
async def predict(data: InputData) -> PredictionResponse:
    """Predict domain classification from preprocessed features.

    Args:
        data: InputData with exactly 5 preprocessed features

    Returns:
        PredictionResponse with predicted_label and raw_prediction_encoded

    Raises:
        HTTPException: If model not loaded or feature count incorrect
    """
    # Check if model is loaded
    if not model_loaded or model is None or label_encoder is None:
        logger.error("Model not loaded, cannot make prediction")
        raise HTTPException(
            status_code=503,
            detail="Model not loaded. Check server logs."
        )

    try:
        # Validate feature count
        if len(data.features) != EXPECTED_FEATURES:
            raise ValueError(
                f"Expected {EXPECTED_FEATURES} features, got {len(data.features)}"
            )

        # Validate all features are float/convertible
        features_array = np.array(data.features, dtype=float)

        # Check for NaN or Inf
        if not np.all(np.isfinite(features_array)):
            raise ValueError("Features contain NaN or Inf values")

        # Reshape for prediction (1 sample, 5 features)
        input_array = features_array.reshape(1, -1)

        logger.debug(f"Input features: {data.features}")

        # Make prediction
        prediction_encoded = model.predict(input_array)[0]
        logger.debug(f"Raw prediction: {prediction_encoded}")

        # Ensure prediction is valid
        if prediction_encoded not in label_encoder.classes_:
            logger.warning(f"Unexpected prediction value: {prediction_encoded}")

        # Decode prediction to label
        try:
            prediction_label = label_encoder.inverse_transform([prediction_encoded])[0]
        except Exception as e:
            logger.error(f"Error decoding prediction: {e}")
            raise ValueError(f"Could not decode prediction {prediction_encoded}")

        logger.info(f"Prediction: {prediction_label}")

        return PredictionResponse(
            predicted_label=prediction_label,
            raw_prediction_encoded=int(prediction_encoded)
        )

    except ValueError as e:
        logger.error(f"Validation error: {e}")
        raise HTTPException(
            status_code=400,
            detail=str(e)
        )
    except Exception as e:
        logger.error(f"Prediction error: {e}")
        import traceback
        traceback.print_exc()
        raise HTTPException(
            status_code=500,
            detail=f"Prediction failed: {str(e)}"
        )


# ─────────────────────────────────────────────────────────────────
# Development Entry Point
# ─────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import uvicorn

    # For development/testing
    logger.info("Starting Model 4 API in development mode...")
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=7860,  # Default HuggingFace Spaces port
        log_level="info"
    )