"""Base regressor with shared sklearn interface and validation logic."""

from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional

import numpy as np
from sklearn.base import BaseEstimator, RegressorMixin
from sklearn.utils.multiclass import type_of_target
from sklearn.utils.validation import check_array, check_is_fitted, check_X_y


class BaseRegressor(RegressorMixin, BaseEstimator, ABC):
    """
    Abstract base regressor with shared sklearn interface.

    This class provides common functionality for both cloud and on-premise
    regressors, including input validation and the standard sklearn
    fit/predict interface.

    Parameters
    ----------
    model : str, default="nicl-small"
        Model identifier to use for inference. Available models:
        - "nicl-flash": Fastest model, optimized for low latency
        - "nicl-small": Balanced model (default)
        - "nicl-large": Most accurate model, recommended for complex tasks
    strategy : str, optional
        Prompting strategy for group-wise processing:
        - "feature": Groups based on specified features
        - "random": Random assignment to groups
        - "correlation": Groups based on target-correlated features
        - "precomputed_groups": Uses pre-existing group IDs
    memory_optimization : bool, default=True
        Enable server-side memory optimization.
    preprocess : bool, default=True
        Enable server-side data preprocessing.
    n_groups : int, optional
        Number of groups for the prompting strategy.
    column_names : List[str], optional
        Column names corresponding to features in X.
    selected_features : List[str], optional
        Features to use for grouping strategies.
    timeout_s : int, default=900
        Request timeout in seconds.
    metadata : dict, optional
        Optional metadata to include with requests.
    user : str, optional
        Optional user identifier for request tracking.

    Attributes
    ----------
    X_train_ : ndarray of shape (n_samples, n_features)
        Training data stored during fit.
    y_train_ : ndarray of shape (n_samples,)
        Training target values stored during fit.
    last_response_ : dict, optional
        The last response received from the server.
    """

    def __init__(
        self,
        *,
        model: str = "nicl-small",
        strategy: Optional[str] = None,
        memory_optimization: bool = True,
        preprocess: bool = True,
        n_groups: Optional[int] = None,
        column_names: Optional[List[str]] = None,
        selected_features: Optional[List[str]] = None,
        timeout_s: int = 900,
        metadata: Optional[Dict[str, Any]] = None,
        user: Optional[str] = None,
    ):
        self.model = model
        self.strategy = strategy
        self.memory_optimization = memory_optimization
        self.preprocess = preprocess
        self.n_groups = n_groups
        self.column_names = column_names
        self.selected_features = selected_features
        self.timeout_s = timeout_s
        self.metadata = metadata
        self.user = user

    def fit(self, X, y) -> "BaseRegressor":
        """
        Fit the regressor by storing training data.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data.
        y : array-like of shape (n_samples,)
            Target values (continuous).

        Returns
        -------
        self : BaseRegressor
            Returns the instance itself.
        """
        X, y = check_X_y(X, y, dtype=None, ensure_2d=True)

        # Validate target type - reject binary/string targets
        target_type = type_of_target(y)
        if target_type in ("binary", "unknown"):
            raise ValueError(
                f"y_train must be a regression target (continuous), "
                f"got '{target_type}'"
            )

        self.X_train_ = check_array(X, dtype=np.float32, order="C")
        self.y_train_ = check_array(
            y, dtype=np.float64, ensure_2d=False, order="C"
        )

        return self

    def predict(self, X) -> np.ndarray:
        """
        Predict target values for samples in X.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Samples to predict.

        Returns
        -------
        y_pred : ndarray of shape (n_samples,)
            Predicted target values per sample.
        """
        check_is_fitted(self, ("X_train_", "y_train_"))
        X = check_array(X, dtype=np.float32, order="C")

        result = self._remote_predict(X)
        predictions = np.asarray(result["predictions"], dtype=np.float64)

        return predictions

    @abstractmethod
    def _remote_predict(self, X: np.ndarray) -> Dict[str, Any]:
        """
        Execute remote prediction.

        Parameters
        ----------
        X : ndarray of shape (n_samples, n_features)
            Validated test samples.

        Returns
        -------
        dict
            Response containing at minimum:
            - predictions: List[float]
        """
        pass

    def get_metadata(self) -> Optional[Dict[str, Any]]:
        """
        Get metadata from the last server response.

        Returns
        -------
        metadata : dict or None
            Metadata from the last prediction call, or None if not available.
        """
        return getattr(self, "metadata_", None)

    def get_last_response(self) -> Optional[Dict[str, Any]]:
        """
        Get the full last server response.

        Returns
        -------
        response : dict or None
            The complete response from the last prediction call.
        """
        return getattr(self, "last_response_", None)
