import itertools
from collections.abc import Callable, Hashable, Iterator, Mapping
from typing import Any

import numpy as np
import pandas as pd
import pyarrow as pa
import xarray as xr

from . import cftime as cft


Block = dict[Hashable, slice]
Chunks = dict[str, int] | None


# Borrowed from Xarray
def _get_chunk_slicer(
    dim: Hashable, chunk_index: Mapping, chunk_bounds: Mapping
):
    if dim in chunk_index:
        which_chunk = chunk_index[dim]
        return slice(
            chunk_bounds[dim][which_chunk], chunk_bounds[dim][which_chunk + 1]
        )
    return slice(None)


# Adapted from Xarray `map_blocks` implementation.
def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[Block]:
    """Compute block slices for a chunked Dataset."""
    if chunks is not None:
        for_chunking = ds.copy(data=None, deep=False).chunk(chunks)
        chunks = for_chunking.chunks
        del for_chunking
    else:
        chunks = ds.chunks

    assert chunks, "Dataset `ds` must be chunked or `chunks` must be provided."

    # chunks is Dict[str, Tuple[int, ...]] from xarray
    chunk_bounds = {
        dim: np.cumsum((0,) + tuple(c))  # type: ignore[arg-type]
        for dim, c in chunks.items()
    }
    ichunk = {dim: range(len(tuple(c))) for dim, c in chunks.items()}  # type: ignore[arg-type]
    ick, icv = zip(*ichunk.items())  # Makes same order of keys and val.
    chunk_idxs = (dict(zip(ick, i)) for i in itertools.product(*icv))
    blocks = (
        {
            dim: _get_chunk_slicer(dim, chunk_index, chunk_bounds)
            for dim in ds.dims
        }
        for chunk_index in chunk_idxs
    )
    yield from blocks


def explode(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[xr.Dataset]:
    """Explodes a dataset into its chunks."""
    yield from (ds.isel(b) for b in block_slices(ds, chunks=chunks))


def _block_len(block: Block) -> int:
    return int(np.prod([v.stop - v.start for v in block.values()]))


def from_map_batched(
    func: Callable[..., pd.DataFrame],
    *iterables: tuple[Any, ...],
    args: tuple | None = None,
    schema: pa.Schema = None,
    **kwargs: dict[str, Any],
) -> pa.RecordBatchReader:
    """Create a PyArrow RecordBatchReader by mapping a function over iterables.

    This is equivalent to dask's from_map but returns a PyArrow
    RecordBatchReader that can be used with DataFusion. It iterates over
    RecordBatches which are created via the `func` one-at-a-time.

    Args:
      func: Function to apply to each element of the iterables. Currently, the
        function must return a Pandas DataFrame.
      *iterables: Iterable objects to map the function over.
      schema: Optional schema needed for the RecordBatchReader.
      args: Additional positional arguments to pass to func.
      **kwargs: Additional keyword arguments to pass to func.

    Returns:
      A PyArrow RecordBatchReader containing the stream of RecordBatches.
    """
    if args is None:
        args = ()

    def map_batches():
        for items in zip(*iterables):
            df = func(*items, *args, **kwargs)
            yield pa.RecordBatch.from_pandas(df, schema=schema)

    return pa.RecordBatchReader.from_batches(schema, map_batches())


def from_map(
    func: Callable,
    *iterables: tuple[Any, ...],
    args: tuple | None = None,
    **kwargs: dict[str, Any],
) -> pa.Table:
    """Create a PyArrow Table by mapping a function over iterables.

    This is equivalent to dask's from_map but returns a PyArrow Table
    that can be used with DataFusion instead of a Dask DataFrame.

    Args:
      func: Function to apply to each element of the iterables.
      *iterables: Iterable objects to map the function over.
      args: Additional positional arguments to pass to func.
      **kwargs: Additional keyword arguments to pass to func.

    Returns:
      A PyArrow Table containing the concatenated results.
    """
    if args is None:
        args = ()

    # Apply the function to each combination of iterable elements
    results = []
    for items in zip(*iterables) if len(iterables) > 1 else iterables[0]:
        if isinstance(items, tuple):
            result = func(*items, *args, **kwargs)
        else:
            result = func(items, *args, **kwargs)

        # Convert result to PyArrow Table
        if isinstance(result, pd.DataFrame):
            pa_table = pa.Table.from_pandas(result)
        elif isinstance(result, pa.Table):
            pa_table = result
        else:
            # Try to convert to pandas first, then to PyArrow
            try:
                df = pd.DataFrame(result)
                pa_table = pa.Table.from_pandas(df)
            except Exception as e:
                raise ValueError(
                    f"Cannot convert function result to PyArrow Table: {e}"
                )

        results.append(pa_table)

    # Concatenate all results
    if not results:
        raise ValueError("No results to concatenate")

    return pa.concat_tables(results)


def pivot(ds: xr.Dataset) -> pd.DataFrame:
    """Converts an xarray Dataset to a pandas DataFrame."""
    return ds.to_dataframe().reset_index()  # type: ignore[no-any-return]


def dataset_to_record_batch(
    ds: xr.Dataset, schema: pa.Schema
) -> pa.RecordBatch:
    """Convert an xarray Dataset partition to an Arrow RecordBatch.

    Builds the RecordBatch directly from numpy arrays, bypassing the pandas
    round-trip (to_dataframe → reset_index → from_pandas) used by pivot().
    For large partitions this reduces peak memory from ~5× to ~2× the
    partition size.

    Dimension coordinates are broadcast to the full partition shape and
    ravelled. np.broadcast_to() is zero-copy; the ravel() forces one copy
    per coordinate (unavoidable, since broadcast arrays are non-contiguous).
    Data variable arrays are ravelled in-place — a zero-copy view when the
    underlying array is already C-contiguous (the common case for numpy-backed
    xarray datasets).

    Args:
        ds: A partition-sized xarray Dataset (already sliced via isel).
        schema: The Arrow schema for the output, as produced by _parse_schema.
            Column order in the output matches schema field order.

    Returns:
        A RecordBatch with one column per dimension coordinate and data
        variable, in schema order.
    """
    # Use the data variable's dimension order as canonical so coordinate
    # broadcasts and data variable ravels use the same layout. All data
    # variables are validated to share the same dims tuple.
    if ds.data_vars:
        first_var = next(iter(ds.data_vars.values()))
        dim_names = list(first_var.dims)
        shape = first_var.shape
    else:
        dim_names = list(ds.sizes.keys())
        shape = tuple(ds.sizes[d] for d in dim_names)

    arrays = []
    for field in schema:
        name = field.name
        if name in ds.coords and name in ds.dims:
            # Broadcast 1-D coordinate to the full N-D partition shape, then ravel.
            axis = dim_names.index(name)
            coord = ds.coords[name].values
            if cft.is_cftime(coord):
                coord = cft.convert_for_field(coord, field)
            reshape = [1] * len(shape)
            reshape[axis] = coord.shape[0]
            arr = np.broadcast_to(coord.reshape(reshape), shape).ravel()
            arrays.append(pa.array(arr, type=field.type))
        else:
            # Data variable: ravel to 1-D (zero-copy for C-contiguous arrays).
            raw = ds[name].values.ravel()
            if cft.is_cftime(ds[name].values):
                raw = cft.convert_for_field(ds[name].values, field)

            # from_pandas=True maps NaN → Arrow null inside the C++ copy kernel,
            # so SQL aggregates (MAX, MIN, AVG) skip missing values correctly.
            arrays.append(pa.array(raw, type=field.type, from_pandas=True))

    return pa.RecordBatch.from_arrays(arrays, schema=schema)


#: Default number of rows per emitted Arrow RecordBatch.
#: 64 K rows balances DataFusion pipeline depth against per-batch overhead.
DEFAULT_BATCH_SIZE: int = 65_536


def iter_record_batches(
    ds: xr.Dataset,
    schema: pa.Schema,
    batch_size: int = DEFAULT_BATCH_SIZE,
) -> Iterator[pa.RecordBatch]:
    """Yield RecordBatches of at most *batch_size* rows from a partition Dataset.

    Unlike `dataset_to_record_batch`, which materialises the entire
    partition as one batch, this generator emits smaller batches so that
    DataFusion can begin filtering and aggregating before the full partition
    is loaded.  Peak memory per batch is O(batch_size) for coordinate columns
    and O(partition_size) for data-variable columns (which must be loaded in
    full from storage).

    Coordinate values are computed per batch via strided index arithmetic —
    no broadcast array spanning the whole partition is ever allocated.  Data
    variable flat arrays are loaded once (triggering any remote I/O) and then
    sliced as zero-copy views for each batch.

    Args:
        ds: A partition-sized xarray Dataset (already sliced via isel).
        schema: The Arrow schema for the output, as produced by _parse_schema.
        batch_size: Maximum number of rows per yielded RecordBatch.

    Yields:
        RecordBatches in schema column order, covering all rows of the
        partition exactly once.
    """
    if ds.data_vars:
        first_var = next(iter(ds.data_vars.values()))
        dim_names = list(first_var.dims)
        shape = first_var.shape
    else:
        dim_names = list(ds.sizes.keys())
        shape = tuple(ds.sizes[d] for d in dim_names)

    total_rows = int(np.prod(shape))

    # Preload small 1-D coordinate arrays (negligible memory).
    # Convert cftime objects to numeric values matching the schema type.
    coord_values = {}
    for name in dim_names:
        vals = ds.coords[name].values
        if cft.is_cftime(vals):
            coord_values[name] = cft.convert_for_field(vals, schema.field(name))
        else:
            coord_values[name] = vals

    # C-order stride for each dimension: stride[k] = prod(shape[k+1:]).
    # Flat row index i → coordinate index for dim k: (i // stride[k]) % shape[k].
    strides = [int(np.prod(shape[k + 1 :])) for k in range(len(shape))]

    # Load data-variable arrays fully (triggers Dask/Zarr compute once).
    # ravel() is a zero-copy view for C-contiguous arrays.
    data_arrays = {}
    for field in schema:
        if field.name not in ds.dims:
            raw = ds[field.name].values
            if cft.is_cftime(raw):
                data_arrays[field.name] = cft.convert_for_field(raw, field)
            else:
                data_arrays[field.name] = raw.ravel()

    for row_start in range(0, total_rows, batch_size):
        row_end = min(row_start + batch_size, total_rows)
        row_idx = np.arange(row_start, row_end)

        arrays = []
        for field in schema:
            name = field.name
            if name in ds.coords and name in ds.dims:
                k = dim_names.index(name)
                coord_idx = (row_idx // strides[k]) % shape[k]
                arrays.append(
                    pa.array(coord_values[name][coord_idx], type=field.type)
                )
            else:
                arrays.append(
                    pa.array(
                        data_arrays[name][row_start:row_end],
                        type=field.type,
                        from_pandas=True,
                    )
                )

        yield pa.RecordBatch.from_arrays(arrays, schema=schema)


def _parse_schema(ds: xr.Dataset) -> pa.Schema:
    """Extracts a `pa.Schema` from the Dataset, treating dims and data_vars as columns.

    Uses the xarray index type to detect cftime coordinates without
    materializing their data — important for Dask/Zarr-backed datasets
    where .values would trigger eager computation.

    cftime coordinates are mapped to one of two Arrow types:

    * **Gregorian-like calendars** (standard, noleap, all_leap, etc.):
      ``pa.timestamp('us')`` so string-based SQL filters work naturally.
    * **Non-Gregorian calendars** (360_day, julian):
      ``pa.int64()`` with ``xarray:units`` / ``xarray:calendar`` metadata
      on the field, preserving lossless CF-convention encoding.
    """
    columns = []

    for coord_name, coord_var in ds.coords.items():
        # Only include dimension coordinates
        if coord_name in ds.dims:
            if cft.is_cftime_index(ds, coord_name):
                units, calendar = cft.encoding(ds, coord_name)
                columns.append(cft.arrow_field(coord_name, units, calendar))
            else:
                pa_type = pa.from_numpy_dtype(coord_var.dtype)
                columns.append(pa.field(coord_name, pa_type))

    for var_name, var in ds.data_vars.items():
        # Data variables are virtually never cftime, but check dtype as a
        # cheap guard.  Only fall back to _is_cftime (which materializes
        # element 0) when dtype is object.
        if var.dtype == np.dtype("O") and cft.is_cftime(var.values):
            # Rare: a data variable holding cftime objects.  Use same encoding
            # as the first cftime dimension coordinate, or default.
            cal = var.values.ravel()[0].calendar
            columns.append(cft.arrow_field(var_name, cft.DEFAULT_UNITS, cal))
        else:
            pa_type = pa.from_numpy_dtype(var.dtype)
            columns.append(pa.field(var_name, pa_type))

    return pa.schema(columns)


# Type alias for partition metadata: maps dimension name to (min, max, dtype_str) values
PartitionBounds = dict[str, tuple[Any, Any, str]]


def _block_metadata(coord_arrays: dict, block: Block) -> PartitionBounds:
    """Compute min/max coordinate values for a single partition block.

    Args:
        coord_arrays: Pre-materialised coordinate arrays keyed by dimension name
            string.  Hoist this outside any loop to avoid repeated remote I/O
            for Zarr-backed datasets.
        block: A single block slice dict from block_slices().

    Returns:
        Dict mapping dimension name to (min_value, max_value, dtype_str).
        Dimensions with an empty slice are omitted; the Rust pruning logic
        treats missing dimensions conservatively (never prunes on them).
    """
    ranges: PartitionBounds = {}
    for dim, slc in block.items():
        coord_values = coord_arrays[str(dim)][slc]
        if len(coord_values) == 0:
            continue
        # String/object dtypes are not representable as ScalarBound
        # (Int64/Float64/TimestampNanos) and numpy min/max ufuncs do not
        # support them.  Skip so pruning treats the dimension conservatively.
        if coord_values.dtype.kind in ("U", "S", "O"):
            continue

        if cft.is_cftime(coord_values):
            ranges[str(dim)] = cft.partition_bounds(coord_values)
            continue

        # Use actual min/max rather than first/last so that non-monotonic
        # coordinate axes (e.g. descending latitude 90→-90) are handled
        # correctly.  np.min/max work for both numeric and datetime64 arrays.
        min_val = coord_values.min()
        max_val = coord_values.max()

        if isinstance(min_val, (np.datetime64, pd.Timestamp)):
            min_val = int(pd.Timestamp(min_val).value)
            max_val = int(pd.Timestamp(max_val).value)
            ranges[str(dim)] = (min_val, max_val, "timestamp_ns")
        elif hasattr(min_val, "item"):
            min_val = min_val.item()
            max_val = max_val.item()
            dtype = "float64" if isinstance(min_val, float) else "int64"
            ranges[str(dim)] = (min_val, max_val, dtype)
        else:
            dtype = "float64" if isinstance(min_val, float) else "int64"
            ranges[str(dim)] = (min_val, max_val, dtype)
    return ranges


def partition_metadata(
    ds: xr.Dataset, blocks: list[Block]
) -> list[PartitionBounds]:
    """Compute min/max coordinate values for each partition.

    This metadata enables filter pushdown: SQL queries with WHERE clauses
    on dimension columns can prune partitions that can't contain matching rows.

    Args:
        ds: The xarray Dataset containing coordinate values.
        blocks: List of block slices from block_slices().

    Returns:
        List of dicts mapping dimension name to
        (min_value, max_value, dtype_str) tuples.

            - For datetime64, values are nanoseconds since Unix epoch
              (int64), dtype_str is "timestamp_ns"
            - For numeric types, values are Python int or float,
              dtype_str is "int64" or "float64"

    Note:
        If a partition has an empty slice for a dimension, that dimension is
        omitted from the partition's metadata. The Rust pruning logic treats
        missing dimensions conservatively (never prunes on them).
    """
    # Hoist coordinate array reads outside the partition loop.
    # ds.coords[dim].values materializes the full array on every call; doing it
    # N_partitions × N_dims times is wasteful and, for remote Zarr-backed datasets
    # (e.g. ARCO-ERA5 on GCS), may trigger repeated network I/O.
    coord_arrays = {str(dim): ds.coords[dim].values for dim in ds.dims}
    return [_block_metadata(coord_arrays, block) for block in blocks]