# -*- coding: utf-8 -*-
"""
Utility functions for the CHARM SDK.

This module provides internal utility functions used by the CHARM SDK,
including validation, and retry configuration.
"""

from typing import Any, Iterable, Sequence, Tuple

from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

from .errors import CharmError, InvalidRequestError


def _validate_shapes(
    descriptions: Sequence[Sequence[str]],
    ts_array: Sequence[Sequence[Sequence[float]]],
):
    """
    Validate input shapes for the CHARM API (embeddings and forecast).

    Expected format (how to shape your data)
    ----------------------------------------
    descriptions : list of list of str, shape (N, C)
        - N = number of time series (batch size).
        - C = number of channels. Each inner list has C strings (channel names).
        Example: [["engine", "temperature"], ["fan", "speed"]] → N=2, C=2.
    ts_array : list of list of list of float, shape (N, T, C)
        - Each of the N elements is one time series of shape (T, C): T timesteps,
          C values per timestep (one per channel).
        - All N series must have the same T and the same C.
        Example: one series with 10 steps, 2 channels →
        [[[0.1, 0.2], [0.3, 0.4], ..., [1.9, 2.0]]] → N=1, T=10, C=2.

    Checks performed
    ----------------
    - len(descriptions) == len(ts_array) (same N).
    - For each sample i: len(descriptions[i]) == number of columns in ts_array[i] (C).
    - All samples have the same T and same C.
    - T >= 1 and T < 1500 (max timesteps per series).

    Parameters
    ----------
    descriptions : Sequence[Sequence[str]]
        Channel names per series; shape (N, C).
    ts_array : Sequence[Sequence[Sequence[float]]]
        Time series values; shape (N, T, C). Same T and C for all.

    Raises
    ------
    InvalidRequestError
        If shapes are inconsistent or T >= 1500.
    """
    if not isinstance(descriptions, Sequence) or not isinstance(ts_array, Sequence):
        raise InvalidRequestError("descriptions and ts_array must be sequences.")

    if len(descriptions) != len(ts_array):
        raise InvalidRequestError(
            f"Batch mismatch: len(descriptions)={len(descriptions)} != len(ts_array)={len(ts_array)}"
        )
    
    ts_shapes = set([len(ts[0]) for ts in ts_array])
    if len(ts_shapes) > 1:
        raise InvalidRequestError(
            f"Batch inconsistency: The batch contains time series with different time steps: {ts_shapes}"
        )

    T = list(ts_shapes)[0]
    if T >= 1500:
        raise InvalidRequestError(
            f"Time series too long: The maximum allowed time steps is 1500, but found {T}."
        )
    
    ts_shapes = set([len(ts) for ts in ts_array])
    if len(ts_shapes) > 1:
        raise InvalidRequestError(
            f"Batch inconsistency: The batch contains time series with different channel counts: {ts_shapes}"
        )
    
    for i, (desc, ts) in enumerate(zip(descriptions, ts_array)):
        if not isinstance(desc, Sequence) or not isinstance(ts, Sequence):
            raise InvalidRequestError(f"Item {i}: invalid types in batch.")

        if len(ts) == 0 or not isinstance(ts[0], Sequence):
            raise InvalidRequestError(f"Item {i}: ts_array must be shaped (T, C).")

        T = len(ts)
        C_from_ts = len(ts[0])

        # All rows must have same C
        for r, row in enumerate(ts):
            if len(row) != C_from_ts:
                raise InvalidRequestError(
                    f"Item {i}: row {r} has {len(row)} columns, expected {C_from_ts}."
                )

        C_from_desc = len(desc)
        if C_from_desc != C_from_ts:
            raise InvalidRequestError(
                f"Item {i}: channel mismatch; descriptions has {C_from_desc} items, "
                f"but ts_array has C={C_from_ts} columns."
            )

        if T < 1:
            raise InvalidRequestError(f"Item {i}: must have at least one timestep.")


def _chunk(iterable: Sequence[Any], size: int) -> Iterable[Sequence[Any]]:
    """
    Split an iterable into chunks of specified size.
    
    This function is used for client-side batching of large datasets.
    
    Parameters
    ----------
    iterable : Sequence[Any]
        The sequence to split into chunks
    size : int
        The maximum size of each chunk
        
    Yields
    ------
    Sequence[Any]
        Chunks of the original sequence
    """
    for i in range(0, len(iterable), size):
        yield iterable[i : i + size]


def _build_retry_adapter(
    total: int,
    backoff_factor: float,
    status_forcelist: Tuple[int, ...] = (429, 500, 502, 503, 504),
):
    """
    Build a retry adapter for requests.
    
    This function configures automatic retries for HTTP requests with
    exponential backoff.
    
    Parameters
    ----------
    total : int
        Maximum number of retries
    backoff_factor : float
        Exponential backoff factor
    status_forcelist : Tuple[int, ...], optional
        HTTP status codes that should trigger a retry
        
    Returns
    -------
    HTTPAdapter
        Configured adapter for the requests session
    """
    # Allow retries on POST because the endpoint is idempotent for prediction.
    retry = Retry(
        total=total,
        read=total,
        connect=total,
        backoff_factor=backoff_factor,
        status_forcelist=status_forcelist,
        allowed_methods=frozenset(["POST"]),
        raise_on_status=False,
        respect_retry_after_header=True,
    )
    return HTTPAdapter(max_retries=retry)