# -*- coding: utf-8 -*-
"""
CHARM SDK for Time Series Embeddings.

This module provides both synchronous and asynchronous clients for the CHARM time
series embedding API. It offers:

- OpenAI-style API surface (client.embeddings.create, client.embeddings.async_create)
- API key authentication
- Automatic retries with exponential backoff
- Configurable timeouts
- Client-side batching for large datasets
- Progress tracking with tqdm
- Flexible return types (Python lists, NumPy arrays, or PyTorch tensors)
- Async client with HTTP/2 support and concurrent batch processing

Performance Note
----------------
The synchronous client processes batches sequentially, which can be slow for large
datasets. The asynchronous client processes batches concurrently, making it significantly
faster for large datasets. The synchronous client should be used with caution for large
datasets (>100 time series) to avoid timeouts.
"""

from __future__ import annotations

import asyncio
import json
import math
import os
from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple, Union

import numpy as np
import torch
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

# Import httpx for async client
try:
    import httpx
except ImportError:
    httpx = None  # type: ignore

# Optional .env loading
try:
    from dotenv import load_dotenv

    load_dotenv()
except ImportError:
    # python-dotenv is not installed; continue without it
    pass

from .errors import (
    APIError,
    AuthenticationError,
    CharmError,
    InvalidRequestError,
    NetworkError,
    RateLimitError,
    ResponseValidationError,
    ServiceUnavailableError,
    TimeoutError_,
)
from .models import EmbeddingsResponse, ForecastResponse
from .utils import _validate_shapes, _chunk, _build_retry_adapter


def _tensor_from_base64(
    blob_b64: str,
    shape: List[int],
    dtype: str = "float16",
    return_tensors: Literal["list", "np", "torch"] = "list",
) -> Union[List[Any], "np.ndarray", "torch.Tensor"]:
    """
    Reconstruct a tensor from a base64-encoded binary blob.
    
    The forecast endpoint returns JSON with base64-encoded fp16 tensors.
    This helper decodes the base64 string, validates the size, and converts
    to the desired format.
    
    Parameters
    ----------
    blob_b64 : str
        Base64-encoded binary data (fp16 tensor bytes).
    shape : List[int]
        Shape of the tensor.
    dtype : str, default "float16"
        Data type string (typically "float16").
    return_tensors : {"list", "np", "torch"}, default "list"
        Desired output type:
        - "list"  : nested Python lists (float32)
        - "np"    : NumPy ndarray (float32)
        - "torch" : torch.Tensor (float32)
    
    Returns
    -------
    list or np.ndarray or torch.Tensor
        Reconstructed tensor in the requested format.
    
    Raises
    ------
    ResponseValidationError
        If the base64 decoding fails or the binary payload size does
        not match the declared shape.
    """
    import base64
    
    try:
        blob = base64.b64decode(blob_b64)
    except Exception as e:
        raise ResponseValidationError(f"Failed to decode base64 blob: {e}") from e
    
    dtype_str = dtype.lower()
    if dtype_str not in ("float16", "fp16"):
        raise ResponseValidationError(f"Unsupported dtype: {dtype_str}")
    
    n_elems = int(np.prod(shape))
    expected_bytes = n_elems * np.dtype(np.float16).itemsize
    if len(blob) != expected_bytes:
        raise ResponseValidationError(
            f"Binary size mismatch: expected {expected_bytes} bytes for shape {shape}, got {len(blob)}"
        )
    
    # Interpret bytes as fp16 and reshape
    arr_fp16 = np.frombuffer(blob, dtype=np.float16).reshape(shape)
    
    # Upcast to float32 for downstream stability/compat
    arr_f32 = arr_fp16.astype(np.float32)
    
    if return_tensors == "np":
        return arr_f32
    elif return_tensors == "torch":
        return torch.from_numpy(arr_f32)
    else:  # "list"
        return arr_f32.tolist()


def _tensor_from_binary(
    blob: bytes,
    headers: Dict[str, Any],
    return_tensors: Literal["list", "np", "torch"],
) -> Union[List[Any], "np.ndarray", "torch.Tensor"]:
    """
    Reconstruct a tensor from a binary response and headers.

    The server encodes model outputs as:
        - Raw fp16 bytes in the HTTP body
        - Shape in the "X-Shape" header (JSON list)
        - Data type (typically "float16") in the "X-Dtype" header

    This helper:
        - Reads shape/dtype from headers (case-insensitive),
        - Validates the body size,
        - Reconstructs the fp16 array and upcasts to float32 for client use,
        - Returns in the desired format.

    Parameters
    ----------
    blob : bytes
        Raw binary response body (fp16 tensor bytes).
    headers : Dict[str, Any]
        Response headers from the server.
    return_tensors : {"list", "np", "torch"}
        Desired output type:
        - "list"  : nested Python lists (float32)
        - "np"    : NumPy ndarray (float32)
        - "torch" : torch.Tensor (float32)

    Returns
    -------
    list or np.ndarray or torch.Tensor
        Reconstructed embedding tensor in the requested format.

    Raises
    ------
    ResponseValidationError
        If headers are missing or invalid, or the binary payload size does
        not match the declared shape.
    """
    # Normalize header keys to lowercase for case-insensitive lookup
    norm_headers = {str(k).lower(): v for k, v in headers.items()}

    shape_str = norm_headers.get("x-shape")
    if not shape_str:
        raise ResponseValidationError(
            f"Missing X-Shape header in binary response. Got headers: {list(headers.keys())}"
        )

    try:
        shape = json.loads(shape_str)
    except Exception as e:
        raise ResponseValidationError(f"Invalid X-Shape header: {shape_str}") from e

    dtype_str = norm_headers.get("x-dtype", "float16").lower()
    if dtype_str not in ("float16", "fp16"):
        raise ResponseValidationError(f"Unsupported X-Dtype: {dtype_str}")

    n_elems = int(np.prod(shape))
    expected_bytes = n_elems * np.dtype(np.float16).itemsize
    if len(blob) != expected_bytes:
        raise ResponseValidationError(
            f"Binary size mismatch: expected {expected_bytes} bytes for shape {shape}, got {len(blob)}"
        )

    # Interpret bytes as fp16 and reshape
    arr_fp16 = np.frombuffer(blob, dtype=np.float16).reshape(shape)

    # Upcast to float32 for downstream stability/compat
    arr_f32 = arr_fp16.astype(np.float32)

    if return_tensors == "np":
        return arr_f32
    elif return_tensors == "torch":
        return torch.from_numpy(arr_f32)
    else:  # "list"
        return arr_f32.tolist()


class CharmClient:
    """
    Client for the CHARM time-series API (embeddings and forecast/backcast).

    Capabilities
    ------------
    - **Embeddings** (POST /predict): `client.embeddings.create(...)` or
      `await client.embeddings.async_create(...)`. Input: descriptions (N,C),
      ts_array (N,T,C). Output: one vector per series, shape (N, D).
    - **Forecast / backcast** (POST /forecast): `client.prediction.create(...)` or
      `await client.prediction.async_create(...)`. Same input plus target_len
      (positive = forecast, negative = backcast). Output: denormalized_predictions
      shape (N, abs(target_len), C, Q).

    Data shapes
    -----------
    All requests use descriptions (list of C channel names per series) and ts_array
    (list of N series, each T timesteps x C channels). N, T, C must be consistent
    in a batch; T < 1500, C < 1500, N*C*T <= 500000. See the package README
    "Data shapes and API reference" for details.

    Authentication
    --------------
    Priority: 1) api_key argument, 2) env CHARM_API_KEY. Sent as Authorization: Bearer <key>.

    Base URL
    --------
    Priority: 1) base_url argument, 2) env CHARM_BASE_URL.

    Example
    -------
    >>> client = CharmClient(base_url="http://localhost:8080", api_key="...")
    >>> resp = client.embeddings.create(descriptions=[["a","b"]], ts_array=[[[0.1,0.2],[0.3,0.4]]], return_tensors="np")
    >>> resp.embeds.shape  # (1, D)
    """

    def __init__(
        self,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
        timeout: Union[int, float] = 15,
        max_retries: int = 3,
        backoff_factor: float = 0.5,
        organization: Optional[str] = None,
        headers: Optional[Dict[str, str]] = None,
        model: str = "charm-embedding",
        http2: bool = True,
        verify_tls: bool = True,
    ):
        """
        Initialize a CharmClient.

        Parameters
        ----------
        base_url : str, optional
            Base URL of the CHARM embeddings server. If not provided, will
            attempt to use the CHARM_BASE_URL environment variable.
        api_key : str, optional
            API key for authentication. If not provided, will attempt to use
            the CHARM_API_KEY environment variable.
        timeout : int or float, default 15
            Request timeout in seconds.
        max_retries : int, default 3
            Maximum number of retry attempts for transient HTTP errors.
        backoff_factor : float, default 0.5
            Exponential backoff factor between retries.
        organization : str, optional
            Optional organization identifier to send as a header.
        headers : dict, optional
            Additional headers to merge into the default request headers.
        model : str, default "charm-embedding"
            Logical model name; returned in EmbeddingsResponse.
        http2 : bool, default True
            Enable HTTP/2 for the async client (httpx).
        verify_tls : bool, default True
            Whether to verify TLS certificates for HTTPS requests.
        """
        # Resolve base_url and api_key (OpenAI-style ergonomics)
        resolved_base_url = (base_url or os.getenv("CHARM_BASE_URL") or "").rstrip("/")
        if not resolved_base_url:
            raise InvalidRequestError(
                "`base_url` is required (or set CHARM_BASE_URL env var)."
            )
        self.base_url = resolved_base_url
        self.predict_url = f"{self.base_url}/predict"
        self.forecast_url = f"{self.base_url}/forecast"

        self.api_key = api_key or os.getenv("CHARM_API_KEY")
        if not self.api_key:
            raise AuthenticationError(
                "No API key provided. Pass `api_key=` or set CHARM_API_KEY env var."
            )

        # Session with retries for synchronous client
        self.session = requests.Session()
        adapter = _build_retry_adapter(total=max_retries, backoff_factor=backoff_factor)
        self.session.mount("http://", adapter)
        self.session.mount("https://", adapter)

        self.timeout = timeout
        self.model = model
        self.http2 = http2
        self.verify_tls = verify_tls
        self.max_retries = max_retries
        self.backoff_factor = backoff_factor

        # Default headers
        default_headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}",
        }
        if organization:
            default_headers["Charm-Organization"] = organization
        if headers:
            default_headers.update(headers)
        self.default_headers = default_headers

        # Async client handle (httpx.AsyncClient)
        self._async_client: Optional["httpx.AsyncClient"] = None

        # Public sub-namespace (OpenAI-like API surface)
        self.embeddings = _EmbeddingsNamespace(self)
        self.prediction = _PredictionNamespace(self)

    # ---------------------- Low-level HTTP helpers ---------------------- #
    def _post_json(self, url: str, payload: Dict[str, Any]) -> Dict[str, Any]:
        """
        Send a synchronous POST request with a JSON payload and parse JSON response.

        Primarily kept for compatibility or non-binary endpoints (not used for
        the binary /predict in the current setup).

        Raises appropriate CharmError subclasses on HTTP/transport failures.
        """
        try:
            r = self.session.post(
                url,
                headers=self.default_headers,
                json=payload,
                timeout=self.timeout,
            )
        except requests.exceptions.Timeout as e:
            raise TimeoutError_(f"Request to {url} timed out after {self.timeout}s.") from e
        except requests.exceptions.RequestException as e:
            raise NetworkError(f"Network error while calling {url}: {e}") from e

        # Map HTTP status codes to rich errors
        if r.status_code in (401, 403):
            raise AuthenticationError(f"{r.status_code} auth error: {r.text}")
        if r.status_code == 429:
            raise RateLimitError(f"429 Too Many Requests: {r.text}")
        if r.status_code == 400:
            raise InvalidRequestError(f"400 Bad Request: {r.text}")
        if r.status_code == 503:
            raise ServiceUnavailableError(f"503 Service Unavailable: {r.text}")
        if r.status_code in (500, 502, 504):
            raise APIError(f"{r.status_code} Server Error: {r.text}")
        if not (200 <= r.status_code < 300):
            raise APIError(f"Unexpected status {r.status_code}: {r.text}")

        try:
            data = r.json()
        except json.JSONDecodeError as e:
            raise ResponseValidationError(
                f"Response was not valid JSON: {r.text[:300]}..."
            ) from e
        return data

    def _post_binary(self, url: str, payload: Dict[str, Any]) -> Tuple[bytes, Dict[str, Any]]:
        """
        Send a synchronous POST request with JSON payload and expect a binary body.

        This is used for the `/predict` endpoint which returns embedding tensors
        as binary fp16 bytes.

        Parameters
        ----------
        url : str
            Target URL for the request.
        payload : dict
            JSON-serializable request payload.

        Returns
        -------
        (bytes, dict)
            Tuple of (raw response body bytes, response headers).

        Raises
        ------
        CharmError
            Mapped from HTTP status codes and transport errors.
        """
        try:
            r = self.session.post(
                url,
                headers=self.default_headers,
                json=payload,
                timeout=self.timeout,
            )
        except requests.exceptions.Timeout as e:
            raise TimeoutError_(f"Request to {url} timed out after {self.timeout}s.") from e
        except requests.exceptions.RequestException as e:
            raise NetworkError(f"Network error while calling {url}: {e}") from e

        # Map HTTP status codes to rich errors
        if r.status_code in (401, 403):
            raise AuthenticationError(f"{r.status_code} auth error: {r.text}")
        if r.status_code == 429:
            raise RateLimitError(f"429 Too Many Requests: {r.text}")
        if r.status_code == 400:
            raise InvalidRequestError(f"400 Bad Request: {r.text}")
        if r.status_code == 503:
            raise ServiceUnavailableError(f"503 Service Unavailable: {r.text}")
        if r.status_code in (500, 502, 504):
            raise APIError(f"{r.status_code} Server Error: {r.text}")
        if not (200 <= r.status_code < 300):
            raise APIError(f"Unexpected status {r.status_code}: {r.text}")

        return r.content, dict(r.headers)

    async def _async_post_json(self, url: str, payload: Dict[str, Any]) -> Dict[str, Any]:
        """
        Send an asynchronous POST request with JSON payload and parse JSON response.

        This is kept for potential non-binary endpoints; the binary `/predict`
        uses `_async_post_binary` instead.

        Raises appropriate CharmError subclasses on HTTP/transport failures.
        """
        if httpx is None:
            raise ImportError(
                "httpx is required for async methods. Install with `pip install httpx`."
            )

        if self._async_client is None:
            limits = httpx.Limits(max_connections=100, max_keepalive_connections=20)
            self._async_client = httpx.AsyncClient(
                http2=self.http2,
                timeout=self.timeout,
                limits=limits,
                verify=self.verify_tls,
            )

        try:
            r = await self._async_client.post(
                url,
                headers=self.default_headers,
                json=payload,
                timeout=self.timeout,
            )
        except httpx.TimeoutException as e:
            raise TimeoutError_(f"Request to {url} timed out after {self.timeout}s.") from e
        except httpx.HTTPError as e:
            raise NetworkError(f"Network error while calling {url}: {e}") from e

        # Map HTTP status codes to rich errors
        if r.status_code in (401, 403):
            raise AuthenticationError(f"{r.status_code} auth error: {await _safe_text(r)}")
        if r.status_code == 429:
            raise RateLimitError(f"429 Too Many Requests: {await _safe_text(r)}")
        if r.status_code == 400:
            raise InvalidRequestError(f"400 Bad Request: {await _safe_text(r)}")
        if r.status_code == 503:
            raise ServiceUnavailableError(f"503 Service Unavailable: {await _safe_text(r)}")
        if r.status_code in (500, 502, 504):
            raise APIError(f"{r.status_code} Server Error: {await _safe_text(r)}")
        if not (200 <= r.status_code < 300):
            raise APIError(f"Unexpected status {r.status_code}: {await _safe_text(r)}")

        try:
            data = r.json()
        except json.JSONDecodeError as e:
            raise ResponseValidationError(
                f"Response was not valid JSON: {await _safe_text(r)}"
            ) from e
        return data

    async def _async_post_binary(self, url: str, payload: Dict[str, Any]) -> Tuple[bytes, Dict[str, Any]]:
        """
        Send an asynchronous POST request with JSON payload and expect a binary body.

        This is used for the `/predict` endpoint which returns embedding tensors
        as binary fp16 bytes.

        Parameters
        ----------
        url : str
            Target URL for the request.
        payload : dict
            JSON-serializable request payload.

        Returns
        -------
        (bytes, dict)
            Tuple of (raw response body bytes, response headers).

        Raises
        ------
        CharmError
            Mapped from HTTP status codes and transport errors.
        """
        if httpx is None:
            raise ImportError(
                "httpx is required for async methods. Install with `pip install httpx`."
            )

        if self._async_client is None:
            limits = httpx.Limits(max_connections=100, max_keepalive_connections=20)
            self._async_client = httpx.AsyncClient(
                http2=self.http2,
                timeout=self.timeout,
                limits=limits,
                verify=self.verify_tls,
            )

        try:
            r = await self._async_client.post(
                url,
                headers=self.default_headers,
                json=payload,
                timeout=self.timeout,
            )
        except httpx.TimeoutException as e:
            raise TimeoutError_(f"Request to {url} timed out after {self.timeout}s.") from e
        except httpx.HTTPError as e:
            raise NetworkError(f"Network error while calling {url}: {e}") from e

        # Map HTTP status codes to rich errors
        if r.status_code in (401, 403):
            raise AuthenticationError(f"{r.status_code} auth error: {await _safe_text(r)}")
        if r.status_code == 429:
            raise RateLimitError(f"429 Too Many Requests: {await _safe_text(r)}")
        if r.status_code == 400:
            raise InvalidRequestError(f"400 Bad Request: {await _safe_text(r)}")
        if r.status_code == 503:
            raise ServiceUnavailableError(f"503 Service Unavailable: {await _safe_text(r)}")
        if r.status_code in (500, 502, 504):
            raise APIError(f"{r.status_code} Server Error: {await _safe_text(r)}")
        if not (200 <= r.status_code < 300):
            raise APIError(f"Unexpected status {r.status_code}: {await _safe_text(r)}")

        return (await r.aread()), dict(r.headers)

    async def aclose(self) -> None:
        """
        Close the underlying async HTTP client, if any.

        Useful when using the client in long-lived async applications.
        """
        if self._async_client is not None:
            await self._async_client.aclose()
            self._async_client = None

    async def __aenter__(self) -> "CharmClient":
        """Async context manager entry; returns the client instance."""
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Async context manager exit; closes the async HTTP client."""
        await self.aclose()


async def _safe_text(resp: "httpx.Response") -> str:
    """
    Safely extract text from an HTTP response for error messages.

    Attempts to read the response body asynchronously first, then falls back
    to `resp.text` if that fails. Output is truncated to 300 characters to avoid
    excessively long error messages.
    """
    try:
        return (await resp.aread()).decode(errors="ignore")[:300]
    except Exception:
        try:
            return resp.text[:300]
        except Exception:
            return "<unreadable>"


class _EmbeddingsNamespace:
    """
    Namespace for embeddings-related operations attached to a CharmClient.

    Exposed as `client.embeddings` to provide an OpenAI-style API surface:

    - Synchronous: `client.embeddings.create(...)`
    - Asynchronous: `await client.embeddings.async_create(...)`
    """

    def __init__(self, client: CharmClient):
        self._client = client

    # --------------------------- Sync embeddings --------------------------- #
    def create(
        self,
        *,
        descriptions: Sequence[Sequence[str]],
        ts_array: Sequence[Sequence[Sequence[float]]],
        batch_size: Optional[int] = None,
        return_tensors: Literal["list", "np", "torch"] = "list",
        progress: bool = True,
    ) -> EmbeddingsResponse:
        """
        Create embeddings for a batch of multivariate time series (synchronous).

        Sends time series to the CHARM /predict endpoint and returns one embedding
        vector per series. For large N, the SDK splits into batches and processes
        sequentially (use async_create for concurrent batching).

        Input shape (how to shape your data)
        ------------------------------------
        descriptions : list of list of str, shape (N, C)
            Channel names for each sample. N = number of time series, C = channels.
            Example: [["engine", "temperature"], ["fan", "speed"]] for N=2, C=2.
        ts_array : list of list of list of float, shape (N, T, C)
            Values for each sample. Each sample has T timesteps and C channels.
            All samples in the request must have the same T and same C.
            Example: for one 10-step, 2-channel series:
            [[[0.1, 0.2], [0.3, 0.4], ..., [1.9, 2.0]]].
        Constraints: T >= 1, T < 1500, C < 1500, N*C*T <= 500000 per batch.

        Output (what to expect)
        -----------------------
        EmbeddingsResponse.embeds : shape (N, D)
            One embedding vector per time series. D = embedding dimension (model-dependent).
        EmbeddingsResponse.model : str
            Model name (e.g. "charm-embedding").
        EmbeddingsResponse.usage : dict or None
            Optional usage stats from the server.

        Parameters
        ----------
        descriptions : Sequence[Sequence[str]]
            Per-sample channel names, shape (N, C).
        ts_array : Sequence[Sequence[Sequence[float]]]
            Per-sample time series, shape (N, T, C). Same T and C for all.
        batch_size : int, optional
            Max samples per API call; SDK splits and concatenates. Defaults to N.
        return_tensors : {"list","np","torch"}, default "list"
            Type of response.embeds: nested lists, numpy.ndarray, or torch.Tensor.
        progress : bool, default True
            If True, show tqdm progress over batches.

        Returns
        -------
        EmbeddingsResponse
            .embeds has shape (N, D). .model, .usage, .raw.

        Raises
        ------
        InvalidRequestError
            If shapes are inconsistent, T >= 1500, or N*C*T > 500000.
        ResponseValidationError
            If the server response is malformed.
        CharmError
            For HTTP or API errors.
        """
        _validate_shapes(descriptions, ts_array)

        N = len(descriptions)
        if batch_size is None:
            batch_size = N
        if not isinstance(batch_size, int) or batch_size <= 0:
            raise InvalidRequestError("batch_size must be a positive integer.")
        
        C = len(ts_array[0])
        T = len(ts_array[0][0])
        if batch_size * C * T > 500_000:
            raise InvalidRequestError(
                f"Batch too large: The maximum allowed size is 500,000 time points, "
                f"but found {N*C*T} (batch size={N}, C={C}, T={T})."
            )

        first_payload_meta: Dict[str, Any] = {}
        all_embeds_np: Optional[np.ndarray] = None

        # Setup tqdm (lazy import) if requested
        bar = None
        if progress:
            try:
                from tqdm.auto import tqdm  # type: ignore
            except ImportError as e:
                raise InvalidRequestError(
                    "tqdm is required for progress=True. Install with `pip install tqdm`."
                ) from e
            total_batches = math.ceil(N / batch_size)
            bar = tqdm(total=total_batches, unit="batch", desc="Embedding batches")

        try:
            for b_idx, (desc_chunk, ts_chunk) in enumerate(
                zip(_chunk(descriptions, batch_size), _chunk(ts_array, batch_size))
            ):
                payload = {
                    "descriptions": desc_chunk,
                    "ts_array": ts_chunk,
                    "batch_size": len(desc_chunk),
                }

                # Binary response from /predict
                blob, headers = self._client._post_binary(self._client.predict_url, payload)

                # Reconstruct as NumPy first; we convert to desired type later
                batch_arr = _tensor_from_binary(blob, headers, return_tensors="np")
                if not isinstance(batch_arr, np.ndarray):
                    raise ResponseValidationError(
                        "Internal error: _tensor_from_binary did not return ndarray for np mode."
                    )

                if b_idx == 0:
                    all_embeds_np = batch_arr
                else:
                    all_embeds_np = np.concatenate([all_embeds_np, batch_arr], axis=0)  # type: ignore[arg-type]

                if bar is not None:
                    bar.update(1)
        finally:
            if bar is not None:
                bar.close()

        if all_embeds_np is None:
            raise ResponseValidationError("No embeddings returned from server.")

        # Convert to desired return type
        if return_tensors == "np":
            converted: Union[List[List[float]], "np.ndarray", "torch.Tensor"] = all_embeds_np.astype(
                np.float32
            )
        elif return_tensors == "torch":
            converted = torch.from_numpy(all_embeds_np.astype(np.float32))
        else:  # "list"
            converted = all_embeds_np.astype(np.float32).tolist()

        return EmbeddingsResponse(
            object="list",
            model=self._client.model,
            embeds=converted,
            usage=first_payload_meta.get("usage") if first_payload_meta else None,
            raw=first_payload_meta if first_payload_meta else None,
        )

    # -------------------------- Async embeddings -------------------------- #
    async def async_create(
        self,
        *,
        descriptions: Sequence[Sequence[str]],
        ts_array: Sequence[Sequence[Sequence[float]]],
        max_B_per_request: int = 64,
        concurrency_per_call: int = 8,
        return_tensors: Literal["list", "np", "torch"] = "list",
        path: Optional[str] = None,
        progress: bool = True,
    ) -> EmbeddingsResponse:
        """
        Generate embeddings for time series with concurrent batch processing (async).

        Same input/output as embeddings.create; batches are sent concurrently for
        better throughput on large N.

        Input shape (how to shape your data)
        ------------------------------------
        descriptions : list of list of str, shape (N, C)
            Channel names per sample. N = number of series, C = channels.
        ts_array : list of list of list of float, shape (N, T, C)
            Values per sample; each sample (T, C). All samples must share same T and C.
        Constraints: T >= 1, T < 1500, C < 1500, N*C*T <= 500000 per request.

        Output (what to expect)
        -----------------------
        EmbeddingsResponse.embeds : shape (N, D)
            One embedding per series. D = embedding dimension.

        Parameters
        ----------
        descriptions : Sequence[Sequence[str]]
            Channel names, shape (N, C).
        ts_array : Sequence[Sequence[Sequence[float]]]
            Time series, shape (N, T, C). Same T and C for all.
        max_B_per_request : int, default 64
            Max time series per API call (each call respects 500k limit on N*C*T).
        concurrency_per_call : int, default 8
            Max concurrent API calls.
        return_tensors : {"list","np","torch"}, default "list"
            Type of response.embeds.
        path : str, optional
            Override endpoint path. Default uses client's /predict URL.
        progress : bool, default True
            Show tqdm progress over batches.

        Returns
        -------
        EmbeddingsResponse
            .embeds shape (N, D). .model, .usage, .raw.

        Raises
        ------
        InvalidRequestError
            If shapes inconsistent or limits violated.
        ResponseValidationError
            If server response is malformed.
        CharmError
            For HTTP or API errors.
        """
        if httpx is None:
            raise ImportError(
                "httpx is required for async methods. Install with `pip install httpx`."
            )

        _validate_shapes(descriptions, ts_array)

        # Prepare batches
        N = len(descriptions)
        if not isinstance(max_B_per_request, int) or max_B_per_request <= 0:
            raise InvalidRequestError("max_B_per_request must be a positive integer")
        if not isinstance(concurrency_per_call, int) or concurrency_per_call <= 0:
            raise InvalidRequestError("concurrency_per_call must be a positive integer")

        C = len(ts_array[0])
        T = len(ts_array[0][0])
        if max_B_per_request * C * T > 500_000:
            raise InvalidRequestError(
                f"Batch too large: The maximum allowed size is 500,000 time points, "
                f"but found {max_B_per_request*C*T} (batch size={max_B_per_request}, C={C}, T={T})."
            )

        # Create batches (list of (desc_chunk, ts_chunk))
        batches: List[Tuple[Sequence[Sequence[str]], Sequence[Sequence[Sequence[float]]]]] = []
        for desc_chunk, ts_chunk in zip(
            _chunk(descriptions, max_B_per_request), _chunk(ts_array, max_B_per_request)
        ):
            batches.append((desc_chunk, ts_chunk))

        # Setup progress bar
        pbar = None
        if progress:
            try:
                from tqdm.auto import tqdm  # type: ignore
            except ImportError as e:
                raise InvalidRequestError(
                    "tqdm is required for progress=True. Install with `pip install tqdm`."
                ) from e
            pbar = tqdm(total=len(batches), desc="Processing batches", leave=True)

        # Process batches with bounded concurrency
        semaphore = asyncio.Semaphore(concurrency_per_call)

        async def process_batch(batch_idx: int, desc_batch, ts_batch):
            async with semaphore:
                payload = {
                    "descriptions": desc_batch,
                    "ts_array": ts_batch,
                    "batch_size": len(desc_batch),
                }

                url = path or self._client.predict_url
                blob, headers = await self._client._async_post_binary(url, payload)

                batch_arr = _tensor_from_binary(blob, headers, return_tensors="np")
                if not isinstance(batch_arr, np.ndarray):
                    raise ResponseValidationError(
                        "Internal error: _tensor_from_binary did not return ndarray for np mode."
                    )

                if pbar:
                    pbar.update(1)

                return batch_idx, batch_arr

        try:
            # Create tasks for all batches
            tasks = [
                process_batch(i, desc_batch, ts_batch)
                for i, (desc_batch, ts_batch) in enumerate(batches)
            ]

            # Wait for all tasks to complete
            results = await asyncio.gather(*tasks)

            # Sort results by batch index to maintain order
            sorted_results = sorted(results, key=lambda x: x[0])

            all_np_batches: List[np.ndarray] = []
            for batch_idx, batch_arr in sorted_results:
                if not isinstance(batch_arr, np.ndarray):
                    raise ResponseValidationError(
                        f"Internal error: batch {batch_idx} did not return ndarray."
                    )
                all_np_batches.append(batch_arr)

            if not all_np_batches:
                raise ResponseValidationError("No embeddings returned from server.")

            all_embeds_np = np.concatenate(all_np_batches, axis=0)

        finally:
            if pbar:
                pbar.close()

        # Convert to desired return type
        if return_tensors == "np":
            converted: Union[List[List[float]], "np.ndarray", "torch.Tensor"] = all_embeds_np.astype(
                np.float32
            )
        elif return_tensors == "torch":
            converted = torch.from_numpy(all_embeds_np.astype(np.float32))
        else:  # "list"
            converted = all_embeds_np.astype(np.float32).tolist()

        return EmbeddingsResponse(
            object="list",
            model=self._client.model,
            embeds=converted,
            usage=None,
            raw=None,
        )


class _PredictionNamespace:
    """
    Namespace for prediction-related operations attached to a CharmClient.
    
    Exposed as `client.prediction` to provide an OpenAI-style API surface:
    
    - Synchronous: `client.prediction.create(...)`
    - Asynchronous: `await client.prediction.async_create(...)`
    """
    
    def __init__(self, client: CharmClient):
        self._client = client
    
    # --------------------------- Sync forecast --------------------------- #
    def create(
        self,
        *,
        descriptions: Sequence[Sequence[str]],
        ts_array: Sequence[Sequence[Sequence[float]]],
        target_len: int,
        batch_size: Optional[int] = None,
        return_tensors: Literal["list", "np", "torch"] = "list",
        progress: bool = True,
    ) -> ForecastResponse:
        """
        Create forecasts or backcasts for a batch of multivariate time series (synchronous).

        Sends time series to the CHARM /forecast endpoint. Use target_len > 0 for
        forecasting future steps, target_len < 0 for backcasting (reconstructing)
        past steps. Returns quantile predictions and an echo of the input.

        Input shape (how to shape your data)
        ------------------------------------
        descriptions : list of list of str, shape (N, C)
            Channel names per sample. N = number of series, C = channels.
        ts_array : list of list of list of float, shape (N, T, C)
            Values per sample; each sample (T, C). All samples must share same T and C.
        target_len : int (required, non-zero)
            Positive: forecast this many steps ahead (e.g. 10 → next 10 steps).
            Negative: backcast this many steps (e.g. -8 → reconstruct last 8 steps).
        Constraints: T >= 1, T < 1500, C < 1500, N*C*T <= 500000 per batch.

        Output (what to expect)
        -----------------------
        ForecastResponse.denormalized_predictions : shape (N, abs(target_len), C, Q)
            Predictions in original scale. Q = number of quantiles (e.g. 21).
            Primary output for downstream use.
        ForecastResponse.predictions : shape (N, abs(target_len), C, Q)
            Normalized (model-internal) predictions.
        ForecastResponse.data : shape (N, T, C)
            Echo of input ts_array.
        ForecastResponse.mode : "forecast" | "backcast"
            Set from sign of target_len.
        ForecastResponse.target_len : int
            Same as the argument.

        Parameters
        ----------
        descriptions : Sequence[Sequence[str]]
            Channel names, shape (N, C).
        ts_array : Sequence[Sequence[Sequence[float]]]
            Time series, shape (N, T, C). Same T and C for all.
        target_len : int
            Non-zero: positive = forecast steps ahead, negative = backcast steps.
        batch_size : int, optional
            Max samples per API call; SDK splits and concatenates. Defaults to N.
        return_tensors : {"list", "np", "torch"}, default "list"
            Type of prediction/data tensors in the response.
        progress : bool, default True
            If True, show tqdm progress over batches.

        Returns
        -------
        ForecastResponse
            .denormalized_predictions (N, abs(target_len), C, Q), .predictions, .data,
            .target_len, .mode, .raw.

        Raises
        ------
        InvalidRequestError
            If shapes inconsistent, target_len is 0, or limits violated.
        ResponseValidationError
            If server response is malformed.
        CharmError
            For HTTP or API errors.
        """
        _validate_shapes(descriptions, ts_array)
        
        if target_len == 0:
            raise InvalidRequestError("target_len cannot be 0. Use positive values for forecasting or negative values for backcasting.")
        
        N = len(descriptions)
        if batch_size is None:
            batch_size = N
        if not isinstance(batch_size, int) or batch_size <= 0:
            raise InvalidRequestError("batch_size must be a positive integer.")
        
        C = len(ts_array[0])
        T = len(ts_array[0][0])
        if batch_size * C * T > 500_000:
            raise InvalidRequestError(
                f"Batch too large: The maximum allowed size is 500,000 time points, "
                f"but found {batch_size*C*T} (batch size={batch_size}, C={C}, T={T})."
            )
        
        first_payload_meta: Dict[str, Any] = {}
        all_predictions_np: Optional[np.ndarray] = None
        all_data_np: Optional[np.ndarray] = None
        all_denorm_predictions_np: Optional[np.ndarray] = None
        
        # Setup tqdm (lazy import) if requested
        bar = None
        if progress:
            try:
                from tqdm.auto import tqdm  # type: ignore
            except ImportError as e:
                raise InvalidRequestError(
                    "tqdm is required for progress=True. Install with `pip install tqdm`."
                ) from e
            total_batches = math.ceil(N / batch_size)
            bar = tqdm(total=total_batches, unit="batch", desc="Forecast batches")
        
        try:
            for b_idx, (desc_chunk, ts_chunk) in enumerate(
                zip(_chunk(descriptions, batch_size), _chunk(ts_array, batch_size))
            ):
                payload = {
                    "descriptions": desc_chunk,
                    "ts_array": ts_chunk,
                    "target_len": target_len,
                }
                
                # JSON response from /forecast
                response_data = self._client._post_json(self._client.forecast_url, payload)
                
                # Validate response structure
                required_keys = ["predictions", "data", "denormalized_predictions"]
                for key in required_keys:
                    if key not in response_data:
                        raise ResponseValidationError(f"Missing key '{key}' in forecast response")
                    if not isinstance(response_data[key], dict):
                        raise ResponseValidationError(f"Invalid format for '{key}' in forecast response")
                    if "blob" not in response_data[key] or "shape" not in response_data[key]:
                        raise ResponseValidationError(f"Missing 'blob' or 'shape' in '{key}' field")
                
                # Decode each tensor field as NumPy first
                batch_predictions = _tensor_from_base64(
                    response_data["predictions"]["blob"],
                    response_data["predictions"]["shape"],
                    response_data["predictions"].get("dtype", "float16"),
                    return_tensors="np"
                )
                
                batch_data = _tensor_from_base64(
                    response_data["data"]["blob"],
                    response_data["data"]["shape"],
                    response_data["data"].get("dtype", "float16"),
                    return_tensors="np"
                )
                
                batch_denorm_predictions = _tensor_from_base64(
                    response_data["denormalized_predictions"]["blob"],
                    response_data["denormalized_predictions"]["shape"],
                    response_data["denormalized_predictions"].get("dtype", "float16"),
                    return_tensors="np"
                )
                
                # Type checking
                if not isinstance(batch_predictions, np.ndarray):
                    raise ResponseValidationError("Internal error: predictions not ndarray for np mode.")
                if not isinstance(batch_data, np.ndarray):
                    raise ResponseValidationError("Internal error: data not ndarray for np mode.")
                if not isinstance(batch_denorm_predictions, np.ndarray):
                    raise ResponseValidationError("Internal error: denormalized_predictions not ndarray for np mode.")
                
                if b_idx == 0:
                    all_predictions_np = batch_predictions
                    all_data_np = batch_data
                    all_denorm_predictions_np = batch_denorm_predictions
                    first_payload_meta = response_data
                else:
                    all_predictions_np = np.concatenate([all_predictions_np, batch_predictions], axis=0)  # type: ignore[arg-type]
                    all_data_np = np.concatenate([all_data_np, batch_data], axis=0)  # type: ignore[arg-type]
                    all_denorm_predictions_np = np.concatenate([all_denorm_predictions_np, batch_denorm_predictions], axis=0)  # type: ignore[arg-type]
                
                if bar is not None:
                    bar.update(1)
        finally:
            if bar is not None:
                bar.close()
        
        if all_predictions_np is None or all_data_np is None or all_denorm_predictions_np is None:
            raise ResponseValidationError("No forecasts returned from server.")
        
        # Convert to desired return type
        if return_tensors == "np":
            predictions_out: Union[List, "np.ndarray", "torch.Tensor"] = all_predictions_np.astype(np.float32)
            data_out: Union[List, "np.ndarray", "torch.Tensor"] = all_data_np.astype(np.float32)
            denorm_predictions_out: Union[List, "np.ndarray", "torch.Tensor"] = all_denorm_predictions_np.astype(np.float32)
        elif return_tensors == "torch":
            predictions_out = torch.from_numpy(all_predictions_np.astype(np.float32))
            data_out = torch.from_numpy(all_data_np.astype(np.float32))
            denorm_predictions_out = torch.from_numpy(all_denorm_predictions_np.astype(np.float32))
        else:  # "list"
            predictions_out = all_predictions_np.astype(np.float32).tolist()
            data_out = all_data_np.astype(np.float32).tolist()
            denorm_predictions_out = all_denorm_predictions_np.astype(np.float32).tolist()
        
        mode = "backcast" if target_len < 0 else "forecast"
        
        return ForecastResponse(
            object="forecast",
            model=self._client.model,
            predictions=predictions_out,
            data=data_out,
            denormalized_predictions=denorm_predictions_out,
            target_len=target_len,
            mode=mode,
            raw=first_payload_meta if first_payload_meta else None,
        )
    
    # -------------------------- Async forecast -------------------------- #
    async def async_create(
        self,
        *,
        descriptions: Sequence[Sequence[str]],
        ts_array: Sequence[Sequence[Sequence[float]]],
        target_len: int,
        max_B_per_request: int = 64,
        concurrency_per_call: int = 8,
        return_tensors: Literal["list", "np", "torch"] = "list",
        path: Optional[str] = None,
        progress: bool = True,
    ) -> ForecastResponse:
        """
        Generate forecasts or backcasts with concurrent batch processing (async).

        Same input/output as prediction.create; batches are sent concurrently for
        better throughput on large N.

        Input shape (how to shape your data)
        ------------------------------------
        descriptions : list of list of str, shape (N, C)
            Channel names per sample.
        ts_array : list of list of list of float, shape (N, T, C)
            Values per sample; each sample (T, C). Same T and C for all.
        target_len : int (non-zero)
            Positive = forecast steps ahead, negative = backcast steps.
        Constraints: T >= 1, T < 1500, C < 1500, N*C*T <= 500000 per request.

        Output (what to expect)
        -----------------------
        ForecastResponse.denormalized_predictions : shape (N, abs(target_len), C, Q)
            Predictions in original scale. Q = number of quantiles (e.g. 21).
        ForecastResponse.predictions, .data : same shapes as in create().

        Parameters
        ----------
        descriptions : Sequence[Sequence[str]]
            Channel names, shape (N, C).
        ts_array : Sequence[Sequence[Sequence[float]]]
            Time series, shape (N, T, C). Same T and C for all.
        target_len : int
            Non-zero: positive = forecast, negative = backcast.
        max_B_per_request : int, default 64
            Max time series per API call.
        concurrency_per_call : int, default 8
            Max concurrent API calls.
        return_tensors : {"list", "np", "torch"}, default "list"
            Type of response tensors.
        path : str, optional
            Override endpoint path. Default uses client's /forecast URL.
        progress : bool, default True
            Show tqdm progress over batches.

        Returns
        -------
        ForecastResponse
            .denormalized_predictions (N, abs(target_len), C, Q), .predictions, .data,
            .target_len, .mode, .raw.

        Raises
        ------
        InvalidRequestError
            If shapes inconsistent, target_len is 0, or limits violated.
        ResponseValidationError
            If server response is malformed.
        CharmError
            For HTTP or API errors.
        """
        if httpx is None:
            raise ImportError(
                "httpx is required for async methods. Install with `pip install httpx`."
            )
        
        _validate_shapes(descriptions, ts_array)
        
        if target_len == 0:
            raise InvalidRequestError("target_len cannot be 0. Use positive values for forecasting or negative values for backcasting.")
        
        # Prepare batches
        N = len(descriptions)
        if not isinstance(max_B_per_request, int) or max_B_per_request <= 0:
            raise InvalidRequestError("max_B_per_request must be a positive integer")
        if not isinstance(concurrency_per_call, int) or concurrency_per_call <= 0:
            raise InvalidRequestError("concurrency_per_call must be a positive integer")
        
        C = len(ts_array[0])
        T = len(ts_array[0][0])
        if max_B_per_request * C * T > 500_000:
            raise InvalidRequestError(
                f"Batch too large: The maximum allowed size is 500,000 time points, "
                f"but found {max_B_per_request*C*T} (batch size={max_B_per_request}, C={C}, T={T})."
            )
        
        # Create batches (list of (desc_chunk, ts_chunk))
        batches: List[Tuple[Sequence[Sequence[str]], Sequence[Sequence[Sequence[float]]]]] = []
        for desc_chunk, ts_chunk in zip(
            _chunk(descriptions, max_B_per_request), _chunk(ts_array, max_B_per_request)
        ):
            batches.append((desc_chunk, ts_chunk))
        
        # Setup progress bar
        pbar = None
        if progress:
            try:
                from tqdm.auto import tqdm  # type: ignore
            except ImportError as e:
                raise InvalidRequestError(
                    "tqdm is required for progress=True. Install with `pip install tqdm`."
                ) from e
            pbar = tqdm(total=len(batches), desc="Processing batches", leave=True)
        
        # Process batches with bounded concurrency
        semaphore = asyncio.Semaphore(concurrency_per_call)
        
        async def process_batch(batch_idx: int, desc_batch, ts_batch):
            async with semaphore:
                payload = {
                    "descriptions": desc_batch,
                    "ts_array": ts_batch,
                    "target_len": target_len,
                }
                
                url = path or self._client.forecast_url
                response_data = await self._client._async_post_json(url, payload)
                
                # Validate response structure
                required_keys = ["predictions", "data", "denormalized_predictions"]
                for key in required_keys:
                    if key not in response_data:
                        raise ResponseValidationError(f"Missing key '{key}' in forecast response")
                    if not isinstance(response_data[key], dict):
                        raise ResponseValidationError(f"Invalid format for '{key}' in forecast response")
                    if "blob" not in response_data[key] or "shape" not in response_data[key]:
                        raise ResponseValidationError(f"Missing 'blob' or 'shape' in '{key}' field")
                
                # Decode each tensor field as NumPy
                batch_predictions = _tensor_from_base64(
                    response_data["predictions"]["blob"],
                    response_data["predictions"]["shape"],
                    response_data["predictions"].get("dtype", "float16"),
                    return_tensors="np"
                )
                
                batch_data = _tensor_from_base64(
                    response_data["data"]["blob"],
                    response_data["data"]["shape"],
                    response_data["data"].get("dtype", "float16"),
                    return_tensors="np"
                )
                
                batch_denorm_predictions = _tensor_from_base64(
                    response_data["denormalized_predictions"]["blob"],
                    response_data["denormalized_predictions"]["shape"],
                    response_data["denormalized_predictions"].get("dtype", "float16"),
                    return_tensors="np"
                )
                
                # Type checking
                if not isinstance(batch_predictions, np.ndarray):
                    raise ResponseValidationError("Internal error: predictions not ndarray for np mode.")
                if not isinstance(batch_data, np.ndarray):
                    raise ResponseValidationError("Internal error: data not ndarray for np mode.")
                if not isinstance(batch_denorm_predictions, np.ndarray):
                    raise ResponseValidationError("Internal error: denormalized_predictions not ndarray for np mode.")
                
                if pbar:
                    pbar.update(1)
                
                return batch_idx, batch_predictions, batch_data, batch_denorm_predictions, response_data
        
        try:
            # Create tasks for all batches
            tasks = [
                process_batch(i, desc_batch, ts_batch)
                for i, (desc_batch, ts_batch) in enumerate(batches)
            ]
            
            # Wait for all tasks to complete
            results = await asyncio.gather(*tasks)
            
            # Sort results by batch index to maintain order
            sorted_results = sorted(results, key=lambda x: x[0])
            
            all_predictions_batches: List[np.ndarray] = []
            all_data_batches: List[np.ndarray] = []
            all_denorm_predictions_batches: List[np.ndarray] = []
            first_payload_meta: Dict[str, Any] = {}
            
            for batch_idx, batch_preds, batch_data, batch_denorm, response_data in sorted_results:
                if not isinstance(batch_preds, np.ndarray):
                    raise ResponseValidationError(f"Internal error: batch {batch_idx} predictions not ndarray.")
                if not isinstance(batch_data, np.ndarray):
                    raise ResponseValidationError(f"Internal error: batch {batch_idx} data not ndarray.")
                if not isinstance(batch_denorm, np.ndarray):
                    raise ResponseValidationError(f"Internal error: batch {batch_idx} denorm not ndarray.")
                
                all_predictions_batches.append(batch_preds)
                all_data_batches.append(batch_data)
                all_denorm_predictions_batches.append(batch_denorm)
                
                # Store metadata from first batch
                if batch_idx == 0:
                    first_payload_meta = response_data
            
            if not all_predictions_batches:
                raise ResponseValidationError("No forecasts returned from server.")
            
            # Concatenate all batches
            all_predictions_np = np.concatenate(all_predictions_batches, axis=0)
            all_data_np = np.concatenate(all_data_batches, axis=0)
            all_denorm_predictions_np = np.concatenate(all_denorm_predictions_batches, axis=0)
            
        finally:
            if pbar:
                pbar.close()
        
        # Convert to desired return type
        if return_tensors == "np":
            predictions_out: Union[List, "np.ndarray", "torch.Tensor"] = all_predictions_np.astype(np.float32)
            data_out: Union[List, "np.ndarray", "torch.Tensor"] = all_data_np.astype(np.float32)
            denorm_predictions_out: Union[List, "np.ndarray", "torch.Tensor"] = all_denorm_predictions_np.astype(np.float32)
        elif return_tensors == "torch":
            predictions_out = torch.from_numpy(all_predictions_np.astype(np.float32))
            data_out = torch.from_numpy(all_data_np.astype(np.float32))
            denorm_predictions_out = torch.from_numpy(all_denorm_predictions_np.astype(np.float32))
        else:  # "list"
            predictions_out = all_predictions_np.astype(np.float32).tolist()
            data_out = all_data_np.astype(np.float32).tolist()
            denorm_predictions_out = all_denorm_predictions_np.astype(np.float32).tolist()
        
        mode = "backcast" if target_len < 0 else "forecast"
        
        return ForecastResponse(
            object="forecast",
            model=self._client.model,
            predictions=predictions_out,
            data=data_out,
            denormalized_predictions=denorm_predictions_out,
            target_len=target_len,
            mode=mode,
            raw=first_payload_meta if first_payload_meta else None,
        )
