"""Core copy/rechunk logic for zarr arrays."""

from __future__ import annotations

import logging
import pathlib
from collections.abc import Callable, Sequence
from tempfile import TemporaryDirectory

import numpy as np
import zarr
from zarr.codecs import BloscCodec

from .transform import TransformSpec, apply_transform, resolve_sources
from .utils import gen_array, iter_slices, tuple_max

logger = logging.getLogger(__name__)


def _find_primary(
    primary_name: str,
    src_groups: list[zarr.Group],
    dst_name: str,
) -> zarr.Array:
    """Return the first array named *primary_name* found across *src_groups*."""
    for grp in src_groups:
        if primary_name in grp:
            return grp[primary_name]
    sources_desc = ", ".join(str(g.store) for g in src_groups)
    raise KeyError(
        f"Transform: source variable '{primary_name}' not found in any source group "
        f"(needed for output '{dst_name}'). Searched: {sources_desc}"
    )


def _full_slices(array: zarr.Array) -> tuple[slice, ...]:
    """Return a tuple of slices covering the full extent of *array*."""
    return tuple(slice(0, s) for s in array.shape)


def simple_remap(
    region: tuple[slice, ...],
    var_in: zarr.Array,
    var_out: zarr.Array,
) -> None:
    """Copy one chunk-aligned *region* from *var_in* to *var_out*.

    Parameters
    ----------
    region:
        Tuple of slices describing the sub-region to copy.
    var_in:
        Source zarr array.
    var_out:
        Destination zarr array.
    """
    var_out[region] = var_in[region]


def double_remap(
    region: tuple[slice, ...],
    var_in: zarr.Array,
    var_out: zarr.Array,
    tmp_path: pathlib.Path | None = None,
) -> None:
    """Copy *region* from *var_in* to *var_out* via a temporary intermediate array.

    This two-pass strategy avoids holding large chunks in memory when the
    rechunking involves a significant reordering along the last dimension.
    In the first pass the data is written to a temporary zarr with a
    smaller last-dimension chunk size; in the second pass that temporary
    array is read into the final output layout.

    Parameters
    ----------
    region:
        Tuple of slices describing the sub-region to copy.
    var_in:
        Source zarr array.
    var_out:
        Destination zarr array.
    tmp_path:
        Directory in which to create the temporary zarr store.
        Uses the system temp directory when *None*.
    """
    with TemporaryDirectory(dir=tmp_path, suffix=".zarr") as tmpdir:
        tmp_group = zarr.open_group(tmpdir, mode="a")

        # Intermediate chunks: keep leading dims the same as the source,
        # but shrink the last dimension to reduce peak memory usage.
        last_chunk = max(var_in.chunks[-1] // 64, var_out.chunks[-1])
        tmp_chunks = (*var_in.chunks[:-1], last_chunk)

        tmp_var = gen_array(
            tmp_group,
            var_in,
            tmp_chunks,
            compressors=BloscCodec(cname="lz4", clevel=1, shuffle="noshuffle"),
        )

        # Pass 1: source -> tmp  (use source + tmp chunk sizes)
        for sl in iter_slices(region, tuple_max(var_in.chunks, tmp_chunks)):
            simple_remap(sl, var_in, tmp_var)

        # Pass 2: tmp -> output  (use tmp + output chunk sizes)
        for sl in iter_slices(region, tuple_max(tmp_chunks, var_out.chunks)):
            simple_remap(sl, tmp_var, var_out)


def copy_variable(
    name: str,
    src_group: zarr.Group,
    dst_group: zarr.Group,
    rechunk: dict[str, int] | None = None,
    double_remap_dim: str | None = None,
    tmp_path: pathlib.Path | None = None,
    dst_name: str | None = None,
) -> None:
    """Copy a single variable from *src_group* to *dst_group*, optionally rechunking.

    Parameters
    ----------
    name:
        Variable (array) name within *src_group*.
    src_group:
        Source zarr group.
    dst_group:
        Destination zarr group.
    rechunk:
        Mapping of dimension name -> new chunk size.  Requires the source
        array to have ``.attrs['_ARRAY_DIMENSIONS']`` set (xarray/CF
        convention).  Dimensions not listed keep their original chunk size.
    double_remap_dim:
        Name of the dimension along which the double-remap strategy should
        be used.  When *None* (or when the dimension is not in the array),
        the simpler single-pass copy is used.
    tmp_path:
        Directory for temporary files used by the double-remap strategy.
    dst_name:
        Name to use in *dst_group*.  Defaults to *name* when *None*.
    """
    src_var: zarr.Array = src_group[name]
    out_name = dst_name if dst_name is not None else name

    # Determine output chunks ------------------------------------------------
    if rechunk:
        dims: Sequence[str] = src_var.attrs.get("_ARRAY_DIMENSIONS", [])
        new_chunks = list(src_var.chunks)
        for i, dim in enumerate(dims):
            if dim in rechunk:
                new_chunks[i] = rechunk[dim]
        out_chunks = tuple(new_chunks)
    else:
        out_chunks = src_var.chunks

    logger.debug(
        "Creating array '%s': shape=%s chunks=%s -> chunks=%s dtype=%s",
        out_name,
        src_var.shape,
        src_var.chunks,
        out_chunks,
        src_var.dtype,
    )

    # Create destination array -----------------------------------------------
    dst_var: zarr.Array = dst_group.create_array(
        out_name,
        shape=src_var.shape,
        chunks=out_chunks,
        dtype=src_var.dtype,
        compressors=src_var.compressors,
        overwrite=True,
    )
    # Copy metadata; update _ARRAY_DIMENSIONS if we renamed a variable
    attrs = dict(src_var.attrs)
    dst_var.update_attributes(attrs)

    # Choose copy strategy ---------------------------------------------------
    region = _full_slices(src_var)
    dims = src_var.attrs.get("_ARRAY_DIMENSIONS", [])

    if double_remap_dim is not None and double_remap_dim in dims:
        logger.debug("Using double-remap strategy for dim '%s'", double_remap_dim)
        double_remap(region, src_var, dst_var, tmp_path=tmp_path)
    else:
        for sl in iter_slices(region, tuple_max(src_var.chunks, out_chunks)):
            simple_remap(sl, src_var, dst_var)


def copy_variable_transformed(
    dst_name: str,
    entry: str | Callable,
    src_groups: list[zarr.Group] | zarr.Group,
    dst_group: zarr.Group,
    rechunk: dict[str, int] | None = None,
    tmp_path: pathlib.Path | None = None,
) -> None:
    """Write a transformed output variable to *dst_group*.

    The shape, dtype, chunk layout, and dimension metadata are derived from the
    *primary* source array (the first source variable for callables, or the
    renamed source for string entries).  When multiple source groups are given
    the primary array is taken from the first group that contains it.

    Parameters
    ----------
    dst_name:
        Output variable name in *dst_group*.
    entry:
        Either a source variable name (str) for a simple rename/copy, or a
        callable whose parameter names specify the source variables.
    src_groups:
        One or more source zarr groups.  Source variables are looked up across
        all groups in order; the first match is used.
    dst_group:
        Destination zarr group.
    rechunk:
        Dimension -> chunk-size overrides applied to the primary source array's
        chunk layout.
    tmp_path:
        Directory for temporary files (currently unused; reserved for future
        double-remap support with transforms).
    """
    groups: list[zarr.Group] = [src_groups] if isinstance(src_groups, zarr.Group) else src_groups

    # Determine primary source variable (used for shape/dtype/chunks/attrs)
    if isinstance(entry, str):
        primary_name = entry
    else:
        sources = resolve_sources(entry)
        if not sources:
            raise ValueError(
                f"Transform callable for '{dst_name}' has no parameters; "
                "cannot determine source variable shape."
            )
        primary_name = sources[0]

    primary: zarr.Array = _find_primary(primary_name, groups, dst_name)

    # Determine output chunks from primary source ----------------------------
    if rechunk:
        dims: Sequence[str] = primary.attrs.get("_ARRAY_DIMENSIONS", [])
        new_chunks = list(primary.chunks)
        for i, dim in enumerate(dims):
            if dim in rechunk:
                new_chunks[i] = rechunk[dim]
        out_chunks = tuple(new_chunks)
    else:
        out_chunks = primary.chunks

    # For computed transforms we use float64 by default unless the primary
    # array is already a float type; for renames we preserve the source dtype.
    if isinstance(entry, str):
        out_dtype = primary.dtype
    else:
        out_dtype = primary.dtype if np.issubdtype(primary.dtype, np.floating) else np.float64

    logger.debug(
        "Creating array '%s': shape=%s chunks=%s -> chunks=%s dtype=%s (primary source: '%s')",
        dst_name,
        primary.shape,
        primary.chunks,
        out_chunks,
        out_dtype,
        primary_name,
    )

    # Create destination array -----------------------------------------------
    dst_var: zarr.Array = dst_group.create_array(
        dst_name,
        shape=primary.shape,
        chunks=out_chunks,
        dtype=out_dtype,
        compressors=primary.compressors,
        overwrite=True,
    )
    # Copy metadata from primary source; update _ARRAY_DIMENSIONS if needed
    dst_var.update_attributes(dict(primary.attrs))

    # Write chunk by chunk ---------------------------------------------------
    region = _full_slices(primary)
    for sl in iter_slices(region, tuple_max(primary.chunks, out_chunks)):
        apply_transform(sl, entry, groups, dst_var, dst_name)


def copy_group(
    src_group: zarr.Group,
    dst_group: zarr.Group,
    variables: Sequence[str] | None = None,
    rechunk: dict[str, int] | None = None,
    double_remap_dim: str | None = None,
    tmp_path: pathlib.Path | None = None,
    transform: TransformSpec | None = None,
    extra_src_groups: list[zarr.Group] | None = None,
) -> None:
    """Copy arrays from *src_group* to *dst_group*.

    Parameters
    ----------
    src_group:
        Primary source zarr group.
    dst_group:
        Destination zarr group.
    variables:
        List of variable names to copy.  All arrays are copied when *None*.
        Ignored when *transform* is provided (the transform spec determines
        which variables are written).
    rechunk:
        Mapping of dimension name -> new chunk size applied to every variable.
    double_remap_dim:
        Dimension name triggering the double-remap strategy (see
        :func:`double_remap`).  Ignored when *transform* is provided.
    tmp_path:
        Directory for temporary files.
    transform:
        Optional transform spec (see :mod:`zarr_copy.transform`).  When
        provided, output variables are determined by the spec rather than by
        *variables*.
    extra_src_groups:
        Additional source zarr groups searched when resolving source variables
        during a transform.  Only used when *transform* is not *None*.
    """
    # Copy group-level attributes from the primary source
    dst_group.update_attributes(dict(src_group.attrs))

    if transform is not None:
        all_src: list[zarr.Group] = [src_group] + (extra_src_groups or [])
        for dst_name, entry in transform.items():
            logger.info("Transforming -> %s ...", dst_name)
            copy_variable_transformed(
                dst_name,
                entry,
                all_src,
                dst_group,
                rechunk=rechunk,
                tmp_path=tmp_path,
            )
        return

    names = list(variables) if variables else list(src_group.array_keys())

    for name in names:
        if name not in src_group:
            raise KeyError(f"Variable '{name}' not found in source group.")
        logger.info("Copying %s ...", name)
        copy_variable(
            name,
            src_group,
            dst_group,
            rechunk=rechunk,
            double_remap_dim=double_remap_dim,
            tmp_path=tmp_path,
        )
