"""
Custom inference handler for the MLForge test sklearn LinearRegression model.

This handler lets the Hugging Face Inference API serve sklearn/joblib models
that are not natively supported by the built-in pipeline loaders.

HuggingFace loads this file automatically when the repo uses:
  library_name: sklearn
  pipeline_tag: tabular-regression

Expected request format:
    {"inputs": [[5.0, 3.0]]}       → single or multi-sample list of feature vectors
"""

from pathlib import Path
from typing import Any

import joblib
import numpy as np


class EndpointHandler:
    """Custom HF Endpoint handler for sklearn regression models."""

    def __init__(self, path: str = ""):
        model_dir = Path(path)

        # Try the canonical name first, then fall back to any .joblib file
        candidate = model_dir / "sklearn_model.joblib"
        if not candidate.exists():
            matches = list(model_dir.glob("*.joblib"))
            if not matches:
                raise FileNotFoundError(
                    f"No .joblib model file found in {model_dir}"
                )
            candidate = matches[0]

        self.model = joblib.load(candidate)
        self._model_file = candidate.name

    def __call__(self, data: dict[str, Any]) -> list[float]:
        """Run inference.

        Args:
            data: Dict with key ``"inputs"`` containing a list of feature vectors,
                  e.g. ``{"inputs": [[5.0, 3.0], [1.0, 2.0]]}``

        Returns:
            List of predicted float values, one per input row.
        """
        raw = data.get("inputs", data.get("data"))
        if raw is None:
            raise ValueError(
                "Request must contain 'inputs' key with a list of feature vectors."
            )

        X = np.array(raw, dtype=float)
        if X.ndim == 1:
            X = X.reshape(1, -1)

        predictions = self.model.predict(X)
        return predictions.tolist()