"""Low-level API client for Neuralk inference API."""

import base64
import hashlib
import io
import json
import logging
import platform
import sys
import threading
import time
from contextlib import contextmanager
from http import HTTPStatus
from importlib import metadata as importlib_metadata
from typing import Any, Dict, List, Optional

import httpx
import numpy as np
from rich.console import Console
from rich.live import Live
from rich.progress import (
    BarColumn,
    DownloadColumn,
    Progress,
    SpinnerColumn,
    TextColumn,
    TimeElapsedColumn,
    TransferSpeedColumn,
)
from rich.spinner import Spinner
from rich.text import Text
from sklearn.utils.validation import check_array

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

def _is_notebook() -> bool:
    """Check if running inside a Jupyter/IPython notebook."""
    try:
        from IPython import get_ipython

        return get_ipython() is not None
    except ImportError:
        return False


def _make_console() -> Console:
    """Create a Rich console that works in both terminals and notebooks."""
    if _is_notebook():
        return Console(force_jupyter=True)
    return Console(stderr=True)


# Rich console for spinner output
_console = _make_console()


def _is_interactive() -> bool:
    """Check if running in an interactive terminal or notebook."""
    return sys.stderr.isatty() or _is_notebook()


@contextmanager
def _phase_spinner(message: str):
    """Context manager for a spinner during a processing phase."""
    if not _is_interactive() or _is_notebook():
        yield None
        return

    spinner = Spinner("dots", text=Text(f" {message}", style="dim"))
    with Live(spinner, console=_console, refresh_per_second=12, transient=True):
        yield spinner


def _create_upload_progress() -> Progress:
    """Create a Rich Progress instance for upload tracking."""
    return Progress(
        SpinnerColumn(),
        TextColumn("[dim]{task.description}"),
        BarColumn(bar_width=30),
        DownloadColumn(),
        TransferSpeedColumn(),
        TimeElapsedColumn(),
        console=_console,
        transient=True,
    )


def _format_bytes(num_bytes: int) -> str:
    """Format bytes into human-readable string."""
    for unit in ["B", "KB", "MB", "GB"]:
        if abs(num_bytes) < 1024.0:
            return f"{num_bytes:.1f}{unit}"
        num_bytes /= 1024.0
    return f"{num_bytes:.1f}TB"


def _print_status(
    success: bool,
    message: str,
    elapsed_ms: Optional[int] = None,
    details: Optional[str] = None,
):
    """Print a status message with icon."""
    if not _is_interactive():
        return

    if success:
        icon = Text("✓", style="green bold")
        text = Text(f" {message}", style="dim")
        if details:
            text.append(f" ({details})", style="dim cyan")
        if elapsed_ms is not None:
            elapsed_str = (
                f"{elapsed_ms / 1000:.2f}s" if elapsed_ms >= 1000 else f"{elapsed_ms}ms"
            )
            text.append(f" · {elapsed_str}", style="cyan")
    else:
        icon = Text("✗", style="red bold")
        text = Text(f" {message}", style="dim red")

    _console.print(icon + text)


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


class Classifications:
    """Resource for classification inference operations using tar-based protocol."""

    def __init__(self, client: "NeuralkAPI"):
        self._client = client

    def create(
        self,
        *,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_test: np.ndarray,
        dataset_name: str = "dataset",
        model: str = "nicl-small",
        strategy: Optional[str] = None,
        memory_optimization: bool = False,
        preprocess: bool = True,
        n_groups: Optional[int] = None,
        column_names: Optional[List[str]] = None,
        selected_features: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        user: Optional[str] = None,
        problem_type: str = "classification",
    ) -> Dict[str, Any]:
        """
        Create a classification inference request.

        This method sends training data and test data to the Neuralk API
        and returns predictions and probabilities.

        Parameters
        ----------
        X_train : np.ndarray
            Training feature matrix of shape (n_samples, n_features).
        y_train : np.ndarray
            Training labels of shape (n_samples,).
        X_test : np.ndarray
            Test feature matrix of shape (n_test_samples, n_features).
        dataset_name : str, default="dataset"
            Name identifier for the dataset.
        model : str, default="nicl-small"
            Model identifier to use for inference (e.g., "nicl-small", "nicl-large").
        strategy : str, optional
            Prompting strategy: "feature", "random", "correlation", or "precomputed_groups".
        memory_optimization : bool, default=False
            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 for feature-based strategies.
        selected_features : List[str], optional
            Features to use for grouping.
        metadata : dict, optional
            Optional metadata to include with the request.
        user : str, optional
            Optional user identifier for request tracking.

        Returns
        -------
        dict
            Response containing:
            - predictions: np.ndarray or List[int]
            - probabilities: np.ndarray or List[List[float]]
            - request_id: str (if provided by server)
            - credits_consumed: int (cloud mode only)
            - latency_ms: int (if provided by server)
            - Additional metadata fields from server response

        Examples
        --------
        >>> from neuralk import NeuralkAPI
        >>> import numpy as np
        >>>
        >>> client = NeuralkAPI(api_key="nk_live_xxxx")
        >>> X_train = np.array([[1.0, 2.0], [2.0, 3.0], [3.0, 4.0]])
        >>> y_train = np.array([0, 0, 1])
        >>> X_test = np.array([[1.5, 2.5]])
        >>>
        >>> result = client.classifications.create(
        ...     X_train=X_train,
        ...     y_train=y_train,
        ...     X_test=X_test,
        ... )
        >>> print(result["predictions"])
        [0]
        """
        # Build prompter config if strategy is set
        prompter_config = None
        if strategy is not None:
            prompter_config = {
                "strategy": strategy,
                "n_groups": n_groups,
                "column_names": column_names,
                "selected_features": selected_features,
            }

        # Build tar archive and headers
        tar_bytes, headers, prep_stats = self._client._build_tar_and_headers(
            dataset_name=dataset_name,
            X_train=X_train,
            X_test=X_test,
            y_train=y_train,
            model=model,
            prompter_config=prompter_config,
            memory_optimization=memory_optimization,
            preprocess=preprocess,
            metadata=metadata,
            user=user,
            problem_type=problem_type,
        )

        # Make the request
        return self._client._make_request(tar_bytes, headers)


class Regressions:
    """Resource for regression inference operations using tar-based protocol."""

    def __init__(self, client: "NeuralkAPI"):
        self._client = client

    def create(
        self,
        *,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_test: np.ndarray,
        dataset_name: str = "dataset",
        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,
        metadata: Optional[Dict[str, Any]] = None,
        user: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Create a regression inference request.

        This method sends training data and test data to the Neuralk API
        and returns predictions.

        Parameters
        ----------
        X_train : np.ndarray
            Training feature matrix of shape (n_samples, n_features).
        y_train : np.ndarray
            Training target values of shape (n_samples,).
        X_test : np.ndarray
            Test feature matrix of shape (n_test_samples, n_features).
        dataset_name : str, default="dataset"
            Name identifier for the dataset.
        model : str, default="nicl-small"
            Model identifier to use for inference (e.g., "nicl-small", "nicl-large").
        strategy : str, optional
            Prompting strategy: "feature", "random", "correlation", or "precomputed_groups".
        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 for feature-based strategies.
        selected_features : List[str], optional
            Features to use for grouping.
        metadata : dict, optional
            Optional metadata to include with the request.
        user : str, optional
            Optional user identifier for request tracking.

        Returns
        -------
        dict
            Response containing:
            - predictions: np.ndarray or List[float]
            - request_id: str (if provided by server)
            - credits_consumed: int (cloud mode only)
            - latency_ms: int (if provided by server)
            - Additional metadata fields from server response

        Examples
        --------
        >>> from neuralk import NeuralkAPI
        >>> import numpy as np
        >>>
        >>> client = NeuralkAPI(api_key="nk_live_xxxx")
        >>> X_train = np.array([[1.0, 2.0], [2.0, 3.0], [3.0, 4.0]])
        >>> y_train = np.array([1.5, 2.5, 3.5])
        >>> X_test = np.array([[1.5, 2.5]])
        >>>
        >>> result = client.regressions.create(
        ...     X_train=X_train,
        ...     y_train=y_train,
        ...     X_test=X_test,
        ... )
        >>> print(result["predictions"])
        [2.0]
        """
        # Build prompter config if strategy is set
        prompter_config = None
        if strategy is not None:
            prompter_config = {
                "strategy": strategy,
                "n_groups": n_groups,
                "column_names": column_names,
                "selected_features": selected_features,
            }

        # Build tar archive and headers
        tar_bytes, headers, prep_stats = self._client._build_tar_and_headers(
            dataset_name=dataset_name,
            X_train=X_train,
            X_test=X_test,
            y_train=y_train,
            model=model,
            prompter_config=prompter_config,
            memory_optimization=memory_optimization,
            preprocess=preprocess,
            metadata=metadata,
            user=user,
            problem_type="regression",
        )

        # Make the request with upload tracking
        return self._client._make_request(tar_bytes, headers, prep_stats=prep_stats)


class NeuralkAPI:
    """
    Low-level API client for Neuralk inference services.

    This client provides direct access to the Neuralk API using a tar-based
    binary protocol for efficient data transfer. For most users, the higher-level
    `Classifier` class is recommended.

    Parameters
    ----------
    api_key : str, optional
        API key for authentication (e.g., "nk_live_xxxx").
        Required for cloud mode. Optional for on-premise deployments.
    host : str, default="https://api.prediction.neuralk-ai.com"
        Base URL of the Neuralk API or on-premise server.
    timeout_s : int, default=900
        Request timeout in seconds.
    max_retries : int, default=8
        Maximum number of retry attempts for failed requests (timeouts,
        network errors, and server errors 500/502/503/504).
    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.

    Examples
    --------
    Cloud mode with API key:

    >>> from neuralk import NeuralkAPI
    >>> import numpy as np
    >>>
    >>> client = NeuralkAPI(api_key="nk_live_xxxx")
    >>> result = client.classifications.create(
    ...     X_train=np.array([[1.0, 2.0], [2.0, 3.0], [3.0, 4.0]]),
    ...     y_train=np.array([0, 0, 1]),
    ...     X_test=np.array([[1.5, 2.5]]),
    ... )
    >>> print(result["predictions"])
    [0]

    On-premise mode without API key:

    >>> client = NeuralkAPI(host="http://localhost:8000")
    >>> result = client.classifications.create(
    ...     X_train=X_train,
    ...     y_train=y_train,
    ...     X_test=X_test,
    ... )
    """

    DEFAULT_HOST = "https://api.prediction.neuralk-ai.com/api/v1/inference"

    def __init__(
        self,
        api_key: Optional[str] = None,
        host: Optional[str] = None,
        timeout_s: int = 900,
        max_retries: int = 8,
        api_version: Optional[str] = None,
        default_headers: Optional[Dict[str, str]] = None,
    ):
        self.api_key = api_key
        self.host = (host or self.DEFAULT_HOST).rstrip("/")
        self.timeout_s = timeout_s
        self.max_retries = max_retries
        self.api_version = api_version
        self.default_headers = default_headers
        self._user_provided_host = host is not None
        self._classifications = None
        self._regressions = None

    @property
    def classifications(self) -> Classifications:
        """Access classification inference operations."""
        if self._classifications is None:
            self._classifications = Classifications(self)
        return self._classifications

    @property
    def regressions(self) -> Regressions:
        """Access regression inference operations."""
        if self._regressions is None:
            self._regressions = Regressions(self)
        return self._regressions

    @property
    def _api_base_url(self) -> str:
        """Derive the API base URL from the inference host.

        The default host points to ``/api/v1/inference``.  The Terms
        endpoints live at ``/api/v1/terms``, so we strip the trailing
        ``/inference`` segment when present.  For on-premise hosts that
        do not end with ``/inference`` the host is returned as-is.

        Returns
        -------
        str
            Base URL suitable for non-inference API calls
            (e.g. ``https://api.prediction.neuralk-ai.com/api/v1``).
        """
        if self.host.endswith("/inference"):
            return self.host.rsplit("/inference", 1)[0]
        return self.host

    @property
    def _is_cloud(self) -> bool:
        """True when targeting the Neuralk cloud API (not on-premise)."""
        return self.api_key is not None and not self._user_provided_host

    def get_terms(self) -> Dict[str, Any]:
        """Fetch the current Terms of Service version and URL.

        This endpoint is public and does not require an API key.

        Returns
        -------
        dict
            Response containing ``current_version`` (str) and ``url`` (str).

        Raises
        ------
        NeuralkException
            If the request fails.

        Examples
        --------
        >>> client = NeuralkAPI()
        >>> terms = client.get_terms()
        >>> print(terms["current_version"])
        v1.0
        """
        url = f"{self._api_base_url}/terms"
        headers = {"User-Agent": _get_user_agent()}
        with httpx.Client(timeout=self.timeout_s) as client:
            response = client.get(url, headers=headers)
        if response.status_code != 200:
            self._raise_for_status(response)
        return response.json()

    def accept_terms(self, version: str) -> Dict[str, Any]:
        """Accept the Terms of Service for the authenticated organization.

        Parameters
        ----------
        version : str
            The Terms version to accept (e.g. ``"v1.0"``).
            Obtain the current version from :meth:`get_terms`.

        Returns
        -------
        dict
            Response containing ``status`` and ``version``.

        Raises
        ------
        NeuralkException
            If the request fails (missing API key, invalid version, etc.).

        Examples
        --------
        >>> client = NeuralkAPI(api_key="nk_live_xxxx")
        >>> terms = client.get_terms()
        >>> client.accept_terms(terms["current_version"])
        """
        if not self.api_key:
            raise NeuralkException(
                "An API key is required to accept the Terms of Service.",
                HTTPStatus.UNAUTHORIZED,
                "Set api_key when constructing NeuralkAPI.",
            )
        url = f"{self._api_base_url}/terms/accept"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "User-Agent": _get_user_agent(),
        }
        with httpx.Client(timeout=self.timeout_s) as client:
            response = client.post(
                url, json={"version": version}, headers=headers,
            )
        if response.status_code != 200:
            self._raise_for_status(response)
        return response.json()

    def _build_tar_and_headers(
        self,
        *,
        dataset_name: str,
        X_train: np.ndarray,
        X_test: np.ndarray,
        y_train: np.ndarray,
        model: str,
        prompter_config: Optional[Dict[str, Any]],
        memory_optimization: bool,
        preprocess: bool,
        metadata: Optional[Dict[str, Any]],
        user: Optional[str],
        problem_type: str = "classification",
    ) -> tuple:
        """Build tar archive and HTTP headers for the request."""
        # Use float64 for regression, int64 for classification
        y_dtype = np.float64 if problem_type == "regression" else np.int64

        prep_start = time.perf_counter()

        with _phase_spinner("Preparing data"):
            X_train = check_array(X_train, dtype=np.float32, order="C", ensure_all_finite="allow-nan")
            X_test = check_array(X_test, dtype=np.float32, order="C", ensure_all_finite="allow-nan")
            y_train = check_array(y_train, dtype=y_dtype, ensure_2d=False, order="C")

            # Calculate raw data size before compression
            raw_size = X_train.nbytes + X_test.nbytes + y_train.nbytes

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

        prep_elapsed_ms = int((time.perf_counter() - prep_start) * 1000)

        # Log preparation stats
        compression_ratio = raw_size / compressed_size if compressed_size > 0 else 0
        _print_status(
            True,
            "Data prepared",
            prep_elapsed_ms,
            f"{_format_bytes(compressed_size)}, {compression_ratio:.1f}x compression",
        )

        headers = {**(self.default_headers or {})}
        headers.setdefault("Content-Type", "application/x-tar+zstd")
        headers.setdefault("User-Agent", _get_user_agent())

        # Add API key for authenticated requests
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"

        if self.api_version:
            headers.setdefault("Nicl-Version", self.api_version)

        headers["X-Content-SHA256"] = hashlib.sha256(tar_bytes).hexdigest()

        stats = {
            "raw_size": raw_size,
            "compressed_size": compressed_size,
            "compression_ratio": compression_ratio,
            "prep_time_ms": prep_elapsed_ms,
        }

        return tar_bytes, headers, stats

    def _make_request(
        self,
        tar_bytes: bytes,
        headers: Dict[str, str],
        timeout: Optional[int] = None,
        prep_stats: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Dispatch to streaming (cloud) or direct (on-premise) request.

        Parameters
        ----------
        tar_bytes : bytes
            The tar archive content to send.
        headers : dict
            HTTP headers to include with the request.
        timeout : int, optional
            Request timeout in seconds. Uses instance timeout_s if not specified.
        prep_stats : dict, optional
            Preparation stats from _build_tar_and_headers (prep_time_ms,
            raw_size, compressed_size).

        Returns
        -------
        dict
            Decoded response containing predictions and metadata.
        """
        if self._is_cloud:
            return self._make_request_streaming(
                tar_bytes, headers, timeout, prep_stats,
            )
        return self._make_request_direct(tar_bytes, headers, timeout)

    def _make_request_direct(
        self,
        tar_bytes: bytes,
        headers: Dict[str, str],
        timeout: Optional[int] = None,
    ) -> Dict[str, Any]:
        """Make a direct (non-streaming) request for on-premise deployments.

        Parameters
        ----------
        tar_bytes : bytes
            The tar archive content to send.
        headers : dict
            HTTP headers to include with the request.
        timeout : int, optional
            Request timeout in seconds. Uses instance timeout_s if not specified.

        Returns
        -------
        dict
            Decoded response containing predictions and metadata.
        """
        timeout = timeout or self.timeout_s
        host = self.host
        attempts = 0
        retry_statuses = {500, 502, 503, 504}
        max_retries = self.max_retries
        total_size = len(tar_bytes)

        while attempts < max_retries:
            attempts += 1
            upload_start = time.perf_counter()

            try:
                # Upload and inference phases with progress tracking
                if _is_interactive():
                    inference_start = None

                    def _upload_generator():
                        """Generator that tracks upload progress and signals completion."""
                        nonlocal inference_start
                        sent = 0
                        for i in range(0, total_size, 8192):
                            chunk = tar_bytes[i : i + 8192]
                            yield chunk
                            sent += len(chunk)
                            progress.update(task_id, completed=sent)

                        # Upload complete - log it and switch to inference spinner
                        upload_elapsed_ms = int(
                            (time.perf_counter() - upload_start) * 1000
                        )
                        # Remove progress bar and print status
                        progress.stop()
                        _print_status(
                            True,
                            "Data uploaded",
                            upload_elapsed_ms,
                            _format_bytes(total_size),
                        )
                        # Start inference timing
                        inference_start = time.perf_counter()
                        # Start inference spinner (will be stopped when response arrives)
                        inference_spinner.start()

                    progress = _create_upload_progress()
                    inference_spinner = _console.status(
                        "[dim]Running inference...", spinner="dots"
                    )
                    task_id = progress.add_task("Uploading data", total=total_size)
                    progress.start()

                    try:
                        with httpx.Client(timeout=timeout) as client:
                            response = client.post(
                                f"{host}",
                                content=_upload_generator(),
                                headers={**headers, "Content-Length": str(total_size)},
                            )
                    finally:
                        inference_spinner.stop()
                        progress.stop()

                else:
                    # Non-interactive mode: simple upload without progress
                    inference_start = time.perf_counter()
                    with httpx.Client(timeout=timeout) as client:
                        response = client.post(
                            f"{host}", content=tar_bytes, headers=headers
                        )

                # Check response
                if response.status_code == 200:
                    result = self._decode_response(response)
                    inference_elapsed_ms = int(
                        (time.perf_counter() - inference_start) * 1000
                    ) if inference_start else 0

                    # Server-side latency if available
                    server_latency = result.get("latency_ms")
                    latency_detail = (
                        f"server: {server_latency}ms" if server_latency else None
                    )
                    _print_status(
                        True, "Inference complete", inference_elapsed_ms, latency_detail
                    )
                    return result

                elif response.status_code in retry_statuses and attempts < max_retries:
                    time.sleep(0.5 * attempts)
                    continue
                else:
                    _print_status(False, f"Request failed ({response.status_code})")
                    self._raise_for_status(response)

            except httpx.TimeoutException as e:
                if attempts < max_retries:
                    time.sleep(0.5 * attempts)
                    continue
                _print_status(False, f"Timeout after {timeout}s")
                raise NeuralkException(
                    f"Timeout after {timeout}s (after {attempts} attempts)",
                    HTTPStatus.REQUEST_TIMEOUT,
                    str(e),
                ) from e
            except httpx.HTTPError as e:
                if attempts < max_retries:
                    time.sleep(0.5 * attempts)
                    continue
                _print_status(False, "Network error")
                raise NeuralkException(
                    f"Network error while calling Neuralk service (after {attempts} attempts)",
                    HTTPStatus.SERVICE_UNAVAILABLE,
                    str(e),
                ) from e

        _print_status(False, "Max retries exceeded")
        raise NeuralkException(
            "Unexpected response from Neuralk service",
            HTTPStatus.INTERNAL_SERVER_ERROR,
            f"No successful response after {max_retries} retries.",
        )

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

        result: Dict[str, Any] = {}

        # Handle both response formats (cloud uses predictions/probabilities,
        # on-premise uses predict/predict_proba)
        if "predictions" in payload:
            result["predictions"] = payload["predictions"]
        elif "predict" in payload:
            result["predictions"] = payload["predict"]

        if "probabilities" in payload:
            result["probabilities"] = payload["probabilities"]
        elif "predict_proba" in payload:
            result["probabilities"] = payload["predict_proba"]

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

        # Unpack nested metadata if present
        if "metadata" in metadata and isinstance(metadata["metadata"], dict):
            nested_metadata = metadata.pop("metadata")
            if nested_metadata:
                metadata.update(nested_metadata)

        # Add metadata fields to result at top level for easy access
        for key, value in metadata.items():
            result[key] = value

        # Extract request ID from response headers
        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
            if not isinstance(err, dict):
                err = 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

        # Detect TOS_NOT_ACCEPTED (error code 4030110)
        if err and err.get("code") == 4030110:
            terms_version = None
            terms_url = None
            try:
                terms_info = self.get_terms()
                terms_version = terms_info.get("current_version")
                terms_url = terms_info.get("url")
            except Exception:
                pass

            parts = [
                "Terms of Service not accepted.",
                "You must accept the Terms before making API calls.",
                "",
                "Option 1 — CLI:    neuralk terms accept --api-key <KEY>",
                "Option 2 — Python:  client.accept_terms("
                f'"{terms_version or "<version>"}") ',
            ]
            if terms_url:
                parts.append(f"Option 3 — Web:    {terms_url}")
            message = "\n".join(parts)

            raise NeuralkTermsNotAcceptedError(
                message,
                HTTPStatus(status),
                response.text,
                terms_version=terms_version,
                terms_url=terms_url,
            )

        raise NeuralkException(message, HTTPStatus(status), response.text)

    # ------------------------------------------------------------------
    # Streaming mode (cloud)
    # ------------------------------------------------------------------

    def _make_request_streaming(
        self,
        tar_bytes: bytes,
        headers: Dict[str, str],
        timeout: Optional[int] = None,
        prep_stats: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Make a streaming request to the Neuralk cloud API.

        Sends the tar payload, then consumes an NDJSON stream that
        reports queue position, processing status and the final result.

        Parameters
        ----------
        tar_bytes : bytes
            The tar archive content to send.
        headers : dict
            HTTP headers to include with the request.
        timeout : int, optional
            Request timeout in seconds.
        prep_stats : dict, optional
            Preparation stats from _build_tar_and_headers.

        Returns
        -------
        dict
            Decoded response containing predictions and metadata.
        """
        timeout = timeout or self.timeout_s
        url = f"{self.host}?stream=true"
        total_size = len(tar_bytes)
        attempts = 0
        retry_statuses = {500, 502, 503, 504}
        max_retries = self.max_retries
        telemetry: Dict[str, Any] = {
            "compressed_size": total_size,
            "raw_size": prep_stats.get("raw_size") if prep_stats else None,
            "client_prep_ms": prep_stats.get("prep_time_ms") if prep_stats else None,
            "n_retries": 0,
        }

        while attempts < max_retries:
            attempts += 1
            upload_start = time.perf_counter()
            telemetry["n_retries"] = attempts - 1

            try:
                httpx_timeout = httpx.Timeout(
                    connect=30.0,
                    read=timeout,
                    write=timeout,
                    pool=timeout,
                )

                with httpx.Client(timeout=httpx_timeout) as client:
                    # Upload with progress tracking (terminal only,
                    # Rich Progress widgets don't work in notebooks)
                    if _is_interactive() and not _is_notebook():
                        progress = _create_upload_progress()
                        task_id = progress.add_task(
                            "Uploading data", total=total_size,
                        )
                        progress.start()

                        def _upload_gen():
                            sent = 0
                            for i in range(0, total_size, 8192):
                                chunk = tar_bytes[i : i + 8192]
                                yield chunk
                                sent += len(chunk)
                                progress.update(task_id, completed=sent)
                            progress.stop()
                            upload_ms = int(
                                (time.perf_counter() - upload_start) * 1000
                            )
                            _print_status(
                                True, "Data uploaded", upload_ms,
                                _format_bytes(total_size),
                            )
                            telemetry["client_upload_ms"] = upload_ms

                        try:
                            with client.stream(
                                "POST",
                                url,
                                content=_upload_gen(),
                                headers={
                                    **headers,
                                    "Content-Length": str(total_size),
                                },
                            ) as response:
                                if response.status_code in retry_statuses:
                                    response.read()
                                    progress.stop()
                                    if attempts < max_retries:
                                        time.sleep(0.5 * attempts)
                                        continue
                                    self._raise_for_status(response)

                                if response.status_code != 200:
                                    response.read()
                                    progress.stop()
                                    self._raise_for_status(response)

                                result = self._consume_ndjson_stream(
                                    response, telemetry,
                                )
                        finally:
                            progress.stop()
                    else:
                        # Notebook or non-interactive: simple upload
                        with client.stream(
                            "POST",
                            url,
                            content=tar_bytes,
                            headers={
                                **headers,
                                "Content-Length": str(total_size),
                            },
                        ) as response:
                            if response.status_code in retry_statuses:
                                response.read()
                                if attempts < max_retries:
                                    time.sleep(0.5 * attempts)
                                    continue
                                self._raise_for_status(response)

                            if response.status_code != 200:
                                response.read()
                                self._raise_for_status(response)

                            upload_ms = int(
                                (time.perf_counter() - upload_start) * 1000
                            )
                            _print_status(
                                True, "Data uploaded", upload_ms,
                                _format_bytes(total_size),
                            )
                            telemetry["client_upload_ms"] = upload_ms
                            result = self._consume_ndjson_stream(
                                response, telemetry,
                            )

                # Telemetry fire-and-forget
                telemetry["outcome"] = "success"
                telemetry["client_total_ms"] = int(
                    (time.perf_counter() - upload_start) * 1000
                )
                if "request_id" in result:
                    telemetry["request_id"] = result["request_id"]
                self._report_telemetry(telemetry)

                return result

            except httpx.TimeoutException as e:
                if attempts < max_retries:
                    time.sleep(0.5 * attempts)
                    continue
                _print_status(False, f"Timeout after {timeout}s")
                telemetry["outcome"] = "timeout"
                telemetry["error_message"] = str(e)
                telemetry["client_total_ms"] = int(
                    (time.perf_counter() - upload_start) * 1000
                )
                self._report_telemetry(telemetry)
                raise NeuralkException(
                    f"Timeout after {timeout}s (after {attempts} attempts)",
                    HTTPStatus.REQUEST_TIMEOUT,
                    str(e),
                ) from e

            except httpx.HTTPError as e:
                if attempts < max_retries:
                    time.sleep(0.5 * attempts)
                    continue
                _print_status(False, "Network error")
                telemetry["outcome"] = "error"
                telemetry["error_message"] = str(e)
                telemetry["client_total_ms"] = int(
                    (time.perf_counter() - upload_start) * 1000
                )
                self._report_telemetry(telemetry)
                raise NeuralkException(
                    f"Network error (after {attempts} attempts)",
                    HTTPStatus.SERVICE_UNAVAILABLE,
                    str(e),
                ) from e

            except NeuralkException:
                telemetry["outcome"] = "error"
                telemetry["client_total_ms"] = int(
                    (time.perf_counter() - upload_start) * 1000
                )
                self._report_telemetry(telemetry)
                raise

        _print_status(False, "Max retries exceeded")
        raise NeuralkException(
            "Unexpected response from Neuralk service",
            HTTPStatus.INTERNAL_SERVER_ERROR,
            f"No successful response after {max_retries} retries.",
        )

    def _consume_ndjson_stream(
        self,
        response: httpx.Response,
        telemetry: Dict[str, Any],
    ) -> Dict[str, Any]:
        """Parse NDJSON events from a streaming inference response.

        Updates the Rich UI with queue position and processing status,
        and returns the final result when the ``complete`` event arrives.

        Parameters
        ----------
        response : httpx.Response
            The streaming HTTP response to read lines from.
        telemetry : dict
            Mutable dict that accumulates client-side timing metrics.

        Returns
        -------
        dict
            Decoded result from the ``complete`` event.

        Raises
        ------
        NeuralkException
            If the server sends an ``error`` event or the stream ends
            without a ``complete`` event.
        """
        queue_start: Optional[float] = None

        try:
            for raw_line in response.iter_lines():
                line = raw_line.strip()
                if not line:
                    continue

                try:
                    event = json.loads(line)
                except json.JSONDecodeError:
                    continue

                event_status = event.get("status")

                if event_status == "queued":
                    queue_start = time.perf_counter()
                    pos = event.get("position", 0)
                    eta = event.get("eta_seconds")
                    active = event.get("active_requests", 0)

                    # Build a human-readable queue message
                    parts = []
                    if pos == 0:
                        parts.append("next up")
                    else:
                        parts.append(
                            f"{pos} request{'s' if pos > 1 else ''} ahead"
                        )
                    if active:
                        parts.append(
                            f"{active} running on server"
                        )
                    if eta:
                        parts.append(f"~{eta}s wait")
                    queue_msg = (
                        "Waiting for slot — " + ", ".join(parts)
                    )
                    _print_status(True, queue_msg)

                elif event_status == "processing":
                    if queue_start is not None:
                        queue_ms = int(
                            (time.perf_counter() - queue_start) * 1000
                        )
                        telemetry["client_queue_wait_ms"] = queue_ms
                        _print_status(True, "Slot acquired", queue_ms)
                    else:
                        _print_status(True, "Slot acquired")

                elif event_status == "complete":
                    inference_ms = event.get("inference_ms")
                    if inference_ms is not None:
                        telemetry["client_inference_ms"] = inference_ms
                    server_latency = event.get("latency_ms")
                    latency_detail = (
                        f"server: {server_latency}ms"
                        if server_latency
                        else None
                    )
                    _print_status(
                        True,
                        "Inference complete",
                        inference_ms,
                        latency_detail,
                    )
                    return self._decode_streaming_complete(event)

                elif event_status == "heartbeat":
                    # Keepalive from server — ignore silently
                    continue

                elif event_status == "error":
                    error_msg = event.get("error", "Unknown streaming error")
                    _print_status(False, f"Error: {error_msg}")
                    telemetry["outcome"] = "error"
                    telemetry["error_message"] = error_msg[:2000]
                    raise NeuralkException(
                        error_msg,
                        HTTPStatus.BAD_GATEWAY,
                        json.dumps(event),
                    )

        finally:
            pass

        raise NeuralkException(
            "Stream ended without a complete event",
            HTTPStatus.BAD_GATEWAY,
            "No complete event received in NDJSON stream.",
        )

    def _decode_streaming_complete(
        self, event: Dict[str, Any],
    ) -> Dict[str, Any]:
        """Decode the ``complete`` event from an NDJSON stream.

        Parameters
        ----------
        event : dict
            The parsed JSON event with status ``complete``.

        Returns
        -------
        dict
            Result dict compatible with ``_decode_response`` output.
        """
        fmt = event.get("format")

        if fmt == "tar+zstd":
            # Binary payload encoded as base64
            raw_b64 = event.get("data", "")
            raw_bytes = base64.b64decode(raw_b64)
            payload = extract_tar(io.BytesIO(raw_bytes), load_cloudpickle=False)

            result: Dict[str, Any] = {}
            if "predictions" in payload:
                result["predictions"] = payload["predictions"]
            elif "predict" in payload:
                result["predictions"] = payload["predict"]

            if "probabilities" in payload:
                result["probabilities"] = payload["probabilities"]
            elif "predict_proba" in payload:
                result["probabilities"] = payload["predict_proba"]

            # Metadata from payload
            meta = {
                k: v
                for k, v in payload.items()
                if k not in {
                    "predict", "predict_proba", "predictions", "probabilities",
                }
            }
            if "metadata" in meta and isinstance(meta["metadata"], dict):
                nested = meta.pop("metadata")
                if nested:
                    meta.update(nested)
            result.update(meta)
        else:
            # JSON-encoded predictions
            result = {}
            if "predictions" in event:
                result["predictions"] = event["predictions"]
            if "probabilities" in event:
                result["probabilities"] = event["probabilities"]

        # Copy streaming metadata into result
        for key in (
            "request_id",
            "credits_consumed",
            "latency_ms",
            "inference_ms",
            "queue_wait_ms",
        ):
            if key in event:
                result[key] = event[key]

        return result

    def _report_telemetry(self, stats: Dict[str, Any]) -> None:
        """Send client telemetry to the backend (fire-and-forget).

        Runs in a daemon thread so it never blocks the caller.

        Parameters
        ----------
        stats : dict
            Telemetry payload with timing metrics and outcome.
        """
        if not self.api_key or not stats.get("request_id"):
            return

        try:
            sdk_version = importlib_metadata.version("neuralk")
        except Exception:
            sdk_version = "unknown"

        payload = {
            "request_id": stats.get("request_id"),
            "sdk_version": sdk_version,
            "python_version": platform.python_version(),
            "client_total_ms": stats.get("client_total_ms", 0),
            "client_prep_ms": stats.get("client_prep_ms"),
            "client_upload_ms": stats.get("client_upload_ms"),
            "client_queue_wait_ms": stats.get("client_queue_wait_ms"),
            "client_inference_ms": stats.get("client_inference_ms"),
            "compressed_size": stats.get("compressed_size"),
            "raw_size": stats.get("raw_size"),
            "outcome": stats.get("outcome", "success"),
            "error_message": stats.get("error_message"),
            "n_retries": stats.get("n_retries", 0),
        }

        url = f"{self._api_base_url}/inference/telemetry"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "User-Agent": _get_user_agent(),
        }

        def _send() -> None:
            try:
                with httpx.Client(timeout=10) as client:
                    client.post(url, json=payload, headers=headers)
            except Exception:
                pass  # fire-and-forget

        thread = threading.Thread(target=_send, daemon=True)
        thread.start()
