import hashlib
import io
import logging
import platform
import time
from http import HTTPStatus
from importlib import metadata as importlib_metadata
from typing import Any, Dict, Optional

import httpx
import numpy as np
from sklearn.base import BaseEstimator, ClassifierMixin
from sklearn.preprocessing import LabelEncoder
from sklearn.utils.validation import check_array, check_is_fitted, check_X_y

from neuralk.exceptions import NeuralkException
from neuralk.utils._remote_utils import create_tar, extract_tar


class OnPremiseClassifier(ClassifierMixin, BaseEstimator):
    """
    Sklearn-style estimator proxying fit/predict calls to an on-premise NICL server.

    This classifier connects to a local or self-hosted NICL (Neural In-Context Learning)
    server and provides a scikit-learn compatible interface for classification tasks.
    It's designed for users who have deployed NICL on their own infrastructure.

    Parameters
    ----------
    host : str
        Base URL of the NICL server (e.g., "http://localhost:8000").
        This is a required parameter.
    dataset_name : str, default="dataset"
        Name identifier for the dataset used in API requests.
    model : str, default="nicl-small"
        Model identifier/path to use for inference (e.g., "nicl-small", "nicl-large").
    timeout_s : int, default=900
        Request timeout in seconds for API calls.
    strategy : str, default=None
        The prompting strategy to use for group-wise processing when data exceeds
        GPU memory capacity. If None (default), the system automatically calculates
        optimal grouping. Available strategies:

        - ``"feature"``: Groups samples based on manually specified features/columns.
          Requires ``column_names`` and ``selected_features`` parameters. Best when
          you have domain knowledge about which features are most relevant for grouping.

        - ``"random"``: Randomly assigns samples to groups. Guarantees perfectly even
          group sizes and is the fastest strategy. Ignores data structure.

        - ``"correlation"``: Automatically selects features most correlated with the
          target variable and groups samples based on quantile distribution of those
          features. Data-driven approach that respects natural data clustering but may
          create uneven group sizes.

        - ``"precomputed_groups"``: Uses pre-existing group IDs from your data.
          Requires ``column_names`` and ``selected_features`` (with exactly one
          feature name). Best when you have already computed optimal group assignments
          externally and want to use them directly.
    memory_optimization: bool, default=True
        Enables internal server-side mechanisms to reduce memory usage and to support
        datasets with an undefined number of columns.In rare cases, these optimizations
        may alter model performance for certain datasets. If the results are unexpected, set this
        flag to False to disable the optimization.
    n_groups : int, default=None
        The number of groups to create for the selected prompting strategy.
    column_names : list[str], default=None
        List of column names corresponding to the features in X, in order.
        Required when using ``strategy="feature"`` or ``strategy="precomputed_groups"``
        with numpy array input. This maps the array columns (which are indexed by
        position) to meaningful names, allowing ``selected_features`` to reference
        columns by name. The length must match the number of features in X.
    selected_features : list[str], default=None
        List of column names to use for grouping when using ``strategy="feature"``,
        or the single column name containing group IDs when using
        ``strategy="precomputed_groups"``. These names must correspond to entries
        in ``column_names``. For ``strategy="feature"``, multiple features can be
        specified.
    metadata : dict, optional
        Optional metadata dictionary to include with requests.
    user : str, optional
        Optional user identifier for request tracking.
    api_version : str, optional
        Optional API version string to send as 'Nicl-Version' header.
    default_headers : dict, optional
        Optional default headers to include with every request (e.g., authentication).

    Attributes
    ----------
    classes_ : ndarray of shape (n_classes,)
        The unique class labels found during fit.
    X_train_ : ndarray of shape (n_samples, n_features)
        Training data stored during fit.
    y_train_ : ndarray of shape (n_samples,)
        Training labels stored during fit.
    last_response_ : dict, optional
        The last response received from the remote server (set after predict/predict_proba).
    metadata_ : dict, optional
        The last metadata received from the remote server (set after predict/predict_proba).

    Examples
    --------
    >>> from neuralk import OnPremiseClassifier
    >>> import numpy as np
    >>>
    >>> # Initialize the classifier with your NICL server URL
    >>> clf = OnPremiseClassifier(
    ...     host="http://localhost:8000",
    ...     model="nicl-small",
    ...     timeout_s=300
    ... )
    >>>
    >>> # Generate some training data
    >>> X_train = np.random.randn(100, 10).astype(np.float32)
    >>> y_train = np.random.randint(0, 2, 100).astype(np.int64)
    >>>
    >>> # Fit the classifier (stores training data)
    >>> clf.fit(X_train, y_train)
    >>>
    >>> # Make predictions
    >>> X_test = np.random.randn(10, 10).astype(np.float32)
    >>> predictions = clf.predict(X_test)
    >>> probabilities = clf.predict_proba(X_test)

    Notes
    -----
    - The classifier requires numeric input data (numpy arrays with float32 dtype).
    - Training labels must be integers (int64 dtype).
    - The fit method only stores the training data; actual model training happens
      on the remote server during predict/predict_proba calls.
    - For advanced use cases, consider using the low-level helpers in
      ``neuralk.model.classify`` directly.
    """

    def __init__(
        self,
        *,
        host: str,
        dataset_name: str = "dataset",
        model: str = "nicl-small",
        timeout_s: int = 900,
        strategy: str | None = None,
        memory_optimization: bool = True,
        n_groups: int | None = None,
        column_names: list[str] | None = None,
        selected_features: list[str] | None = None,
        metadata: Optional[Dict[str, Any]] = None,
        user: Optional[str] = None,
        api_version: Optional[str] = None,
        default_headers: Optional[Dict[str, str]] = None,
    ) -> None:
        self.host = host
        self.dataset_name = dataset_name
        self.model = model
        self.timeout_s = timeout_s
        self.strategy = strategy
        self.memory_optimization = memory_optimization
        self.n_groups = n_groups
        self.column_names = column_names
        self.selected_features = selected_features
        self.metadata = metadata
        self.user = user
        self.api_version = api_version
        self.default_headers = default_headers

    def fit(self, X, y) -> "OnPremiseClassifier":
        """Fit the classifier 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. Can be numeric or string labels.

        Returns
        -------
        self : OnPremiseClassifier
            Returns the instance itself.
        """
        # Validate input data
        X, y = check_X_y(X, y, dtype=None, ensure_2d=True)

        # Handle label encoding for non-numeric labels
        if not np.issubdtype(y.dtype, np.integer):
            self.label_encoder_ = LabelEncoder()
            y_encoded = self.label_encoder_.fit_transform(y)
            self.classes_ = self.label_encoder_.classes_
        else:
            y_encoded = y
            self.classes_ = np.unique(y_encoded)
            self.label_encoder_ = None

        # Convert to required dtypes
        self.X_train_ = check_array(X, dtype=np.float32, order="C")
        self.y_train_ = check_array(
            y_encoded, dtype=np.int64, ensure_2d=False, order="C"
        )

        return self

    def predict(self, X) -> np.ndarray:
        """Predict class labels 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 class labels per sample.
        """
        check_is_fitted(self, ("X_train_", "y_train_"))
        X = check_array(X, dtype=np.float32, order="C")
        result = self._fit_predict_remote(X_test=X)
        predictions = result["predict"]

        # Decode labels if we used a label encoder
        if self.label_encoder_ is not None:
            predictions = self.label_encoder_.inverse_transform(predictions)

        return predictions

    def predict_proba(self, X) -> np.ndarray:
        """Predict class probabilities for samples in X.

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

        Returns
        -------
        y_proba : ndarray of shape (n_samples, n_classes)
            Class probabilities for each sample. The columns correspond to
            the classes in sorted order, as they appear in the attribute
            ``classes_``.
        """
        check_is_fitted(self, ("X_train_", "y_train_"))
        X = check_array(X, dtype=np.float32, order="C")
        result = self._fit_predict_remote(X_test=X)
        return result["predict_proba"]

    def _format_metadata_value(self, value: Any) -> str:
        """Format a metadata value for display."""
        if isinstance(value, np.ndarray):
            if value.ndim == 0:
                return str(value.item())
            elif value.size <= 10:
                return str(value.tolist())
            else:
                return (
                    f"array of shape {value.shape}, dtype {value.dtype}\n"
                    f"  min: {value.min():.4f}, max: {value.max():.4f}, mean: {value.mean():.4f}"
                )
        return str(value)

    def print_metadata(self) -> None:
        """Print the metadata from the last server response.

        This method displays all metadata returned by the server during the last
        predict or predict_proba call. Metadata provides insights into the
        model's performance, configuration, or other relevant information.

        Examples
        --------
        >>> clf = OnPremiseClassifier(host="http://localhost:8000")
        >>> clf.fit(X_train, y_train)
        >>> predictions = clf.predict(X_test)
        >>> clf.print_metadata()  # Shows metadata from the server response
        """
        metadata = self.get_metadata()
        if metadata is None:
            print("No metadata available. Call predict() or predict_proba() first.")
            return

        if not metadata:
            print("No metadata returned by server.")
            return

        print("Server Metadata:")
        print("-" * 40)
        for key, value in metadata.items():
            formatted = self._format_metadata_value(value)
            print(f"{key}: {formatted}")
        print("-" * 40)

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

        Returns
        -------
        metadata : dict or None
            Metadata dictionary returned by the server in the last predict/predict_proba call.
            Returns None if no metadata is available (e.g., before calling predict).

        Examples
        --------
        >>> clf = OnPremiseClassifier(host="http://localhost:8000")
        >>> clf.fit(X_train, y_train)
        >>> predictions = clf.predict(X_test)
        >>> metadata = clf.get_metadata()
        >>> if metadata:
        ...     print(f"Model version: {metadata.get('model_version')}")
        """
        return getattr(self, "metadata_", None)

    def _fit_predict_remote(self, X_test) -> Dict[str, Any]:
        X_test = check_array(X_test, dtype=np.float32, order="C")
        return self._call_remote(
            X_train=self.X_train_,
            X_test=X_test,
            y_train=self.y_train_,
        )

    def _call_remote(
        self,
        *,
        X_train,
        X_test,
        y_train,
        dataset_name: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
        user: Optional[str] = None,
        timeout_s: Optional[int] = None,
        extra_headers: Optional[Dict[str, str]] = None,
    ) -> Dict[str, Any]:
        """Execute the remote fit_predict call."""
        # Prepare request data
        dataset_name = dataset_name or self.dataset_name
        metadata = metadata or self.metadata
        if self.strategy is not None:
            prompter_config = {
                "strategy": self.strategy,
                "n_groups": self.n_groups,
                "column_names": self.column_names,
                "selected_features": self.selected_features,
            }

        user = user or self.user
        timeout = timeout_s or self.timeout_s

        # Build tar archive and headers
        tar_bytes, req_headers, n_train, n_test, num_features = (
            self._build_tar_and_headers(
                dataset_name=dataset_name,
                X_train=X_train,
                X_test=X_test,
                y_train=y_train,
                model_path=self.model,
                prompter_config=prompter_config if self.strategy is not None else None,
                memory_optimization=self.memory_optimization,
                metadata=metadata,
                user=user,
                extra_headers=extra_headers,
            )
        )

        # Make HTTP request
        result = self._make_request(tar_bytes, req_headers, timeout)
        self.last_response_ = result
        return result

    def _build_tar_and_headers(
        self,
        *,
        dataset_name: str,
        X_train,
        X_test,
        y_train,
        model_path: str,
        prompter_config: Optional[Dict[str, Any]],
        memory_optimization: bool,
        metadata: Optional[Dict[str, Any]],
        user: Optional[str],
        extra_headers: Optional[Dict[str, str]],
    ):
        """Build tar archive and HTTP headers for the request."""
        X_train = check_array(X_train, dtype=np.float32, order="C")
        X_test = check_array(X_test, dtype=np.float32, order="C")
        y_train = check_array(y_train, dtype=np.int64, ensure_2d=False, order="C")

        archive = create_tar(
            {
                "method": "fit_predict",
                "model_path": model_path,
                "dataset": dataset_name,
                "prompter_config": prompter_config,
                "memory_optimization": memory_optimization,
                "metadata": metadata or {},
                "user": user or "",
            },
            {"X_train": X_train, "X_test": X_test, "y_train": y_train},
        )
        tar_bytes = archive.getvalue()

        base_headers = {**(self.default_headers or {})}
        base_headers.setdefault("Content-Type", "application/x-tar+zstd")
        base_headers.setdefault("User-Agent", self._get_user_agent())
        if self.api_version:
            base_headers.setdefault("Nicl-Version", self.api_version)
        if extra_headers:
            base_headers.update(extra_headers)
        base_headers["X-Content-SHA256"] = hashlib.sha256(tar_bytes).hexdigest()

        try:
            n_train, num_features = X_train.shape
            n_test = X_test.shape[0]
        except Exception:
            n_train, n_test, num_features = len(X_train), len(X_test), None

        return tar_bytes, base_headers, n_train, n_test, num_features

    def _get_user_agent(self) -> str:
        """Get the User-Agent string for requests."""
        try:
            pkg_version = importlib_metadata.version("neuralk_inference")
        except Exception:
            pkg_version = "0"
        return (
            f"neuralk-inference-client/{pkg_version} httpx/{httpx.__version__} "
            f"python/{platform.python_version()}"
        )

    def _make_request(
        self, tar_bytes: bytes, headers: Dict[str, str], timeout: int
    ) -> Dict[str, Any]:
        """Make the HTTP request to the remote server."""
        host = self.host
        attempts = 0
        retry_statuses = {500, 502, 503, 504}

        while attempts < 3:
            attempts += 1
            start = time.perf_counter()

            try:
                with httpx.Client(timeout=timeout) as client:
                    response = client.post(
                        f"{host}", content=tar_bytes, headers=headers
                    )

                    if response.status_code == 200:
                        elapsed_ms = int((time.perf_counter() - start) * 1000)
                        logging.getLogger("neuralk_nicl_client").info(
                            "POST /fit_predict 200 in %dms", elapsed_ms
                        )
                        return self._decode_response(response)
                    elif response.status_code in retry_statuses and attempts < 3:
                        logging.getLogger("neuralk_nicl_client").warning(
                            "POST /fit_predict %s; retrying %d/3",
                            response.status_code,
                            attempts + 1,
                        )
                        time.sleep(0.5 * attempts)
                        continue
                    else:
                        self._raise_for_status(response)

            except httpx.TimeoutException as e:
                raise NeuralkException(
                    f"Timeout after {timeout}s",
                    HTTPStatus.REQUEST_TIMEOUT,
                    str(e),
                ) from e
            except httpx.HTTPError as e:
                raise NeuralkException(
                    "Network error while calling NICL service",
                    HTTPStatus.SERVICE_UNAVAILABLE,
                    str(e),
                ) from e

        raise NeuralkException(
            "Unexpected response from NICL service",
            HTTPStatus.INTERNAL_SERVER_ERROR,
            "No successful response after retries.",
        )

    def _decode_response(self, response: httpx.Response) -> Dict[str, Any]:
        """Decode the response from the server."""
        response_bytes = response.content
        payload = extract_tar(io.BytesIO(response_bytes), load_cloudpickle=False)
        result: Dict[str, Any] = {
            "predict": payload["predict"],
            "predict_proba": payload["predict_proba"],
        }

        # Extract metadata: everything except predict and predict_proba
        metadata = {
            key: value
            for key, value in payload.items()
            if key not in ["predict", "predict_proba"]
        }

        # If metadata comes as a nested dict under "metadata" key (from metadata.json),
        # extract and merge it with any other metadata keys
        if "metadata" in metadata and isinstance(metadata["metadata"], dict):
            # Get the nested metadata dict
            nested_metadata = metadata.pop("metadata")
            # Merge nested metadata into top level, preserving any other keys
            if nested_metadata:  # Only merge if not empty
                metadata.update(nested_metadata)

        # Store metadata (keep empty dict to distinguish from None/not available)
        self.metadata_ = metadata

        request_id = response.headers.get("x-request-id")
        if request_id:
            result["request_id"] = request_id
        return result

    def _raise_for_status(self, response: httpx.Response) -> None:
        """Raise an appropriate exception for HTTP error responses."""
        try:
            data = response.json()
            err = data.get("error") if isinstance(data, dict) else None
        except Exception:
            err = None
        message = response.text
        if err:
            rid = err.get("request_id")
            message = f"{err.get('code')}: {err.get('message')}"
            if rid:
                message += f" (request_id={rid})"

        status = response.status_code
        if status in (401, 403):
            raise NeuralkException(message, HTTPStatus(status), response.text)
        elif 400 <= status < 500:
            raise NeuralkException(message, HTTPStatus(status), response.text)
        else:
            raise NeuralkException(message, HTTPStatus(status), response.text)
