"""Transform specification: rename and/or compute output variables from source variables."""

from __future__ import annotations

import importlib.util
import inspect
import logging
import pathlib
from collections.abc import Callable
from typing import Union

import numpy as np
import zarr

# A transform spec maps output variable names to either:
#   - a str  -> rename from that source variable (simple copy)
#   - a callable -> compute from source variables whose names match the parameter names
TransformSpec = dict[str, Union[str, Callable]]

logger = logging.getLogger(__name__)


def load_transform_spec(path: pathlib.Path | str) -> TransformSpec:
    """Load a transform spec from a Python file.

    The file must define a module-level variable named ``transform`` that is a
    :data:`TransformSpec` dict.

    Parameters
    ----------
    path:
        Path to the Python config file.

    Returns
    -------
    TransformSpec
        The ``transform`` dict from the config file.

    Raises
    ------
    FileNotFoundError
        If *path* does not exist.
    ValueError
        If the file does not define a ``transform`` variable or it is not a dict.
    """
    path = pathlib.Path(path)
    if not path.exists():
        raise FileNotFoundError(f"Transform config not found: {path}")

    spec = importlib.util.spec_from_file_location("_zarr_copy_transform_config", path)
    if spec is None or spec.loader is None:
        raise ValueError(f"Cannot load Python file: {path}")

    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)  # type: ignore[union-attr]

    if not hasattr(module, "transform"):
        raise ValueError(f"Transform config {path} must define a module-level 'transform' dict.")

    result = module.transform
    if not isinstance(result, dict):
        raise ValueError(f"'transform' in {path} must be a dict, got {type(result).__name__}.")
    logger.debug(
        "Loaded transform spec from '%s': %d output variable(s): %s",
        path,
        len(result),
        list(result),
    )
    return result


def resolve_sources(func: Callable) -> list[str]:
    """Return the source variable names required by *func*.

    The source variable names are inferred from the parameter names of *func*
    using :func:`inspect.signature`.

    Parameters
    ----------
    func:
        A callable whose positional parameter names correspond to source
        variable names in the zarr group.

    Returns
    -------
    list[str]
        Ordered list of source variable names.
    """
    sig = inspect.signature(func)
    sources = [
        name
        for name, param in sig.parameters.items()
        if param.kind
        in (
            inspect.Parameter.POSITIONAL_ONLY,
            inspect.Parameter.POSITIONAL_OR_KEYWORD,
        )
    ]
    logger.debug(
        "Resolved sources for callable '%s': %s", getattr(func, "__name__", repr(func)), sources
    )
    return sources


def _find_var(
    src_name: str,
    src_groups: list[zarr.Group],
    dst_name: str,
) -> zarr.Array:
    """Return the first array named *src_name* found across *src_groups*.

    Raises
    ------
    KeyError
        If *src_name* is not present in any of the groups.
    """
    for grp in src_groups:
        if src_name in grp:
            return grp[src_name]
    sources_desc = ", ".join(str(g.store) for g in src_groups)
    raise KeyError(
        f"Transform: source variable '{src_name}' not found in any source group "
        f"(needed for output '{dst_name}'). Searched: {sources_desc}"
    )


def apply_transform(
    region: tuple[slice, ...],
    entry: str | Callable,
    src_groups: list[zarr.Group] | zarr.Group,
    dst_var: zarr.Array,
    dst_name: str,
) -> None:
    """Apply one transform spec entry for *region* and write the result to *dst_var*.

    Parameters
    ----------
    region:
        Tuple of slices describing the sub-region to process.
    entry:
        Either a source variable name (str) or a callable that computes the
        output from one or more source arrays.
    src_groups:
        One or more source zarr groups.  When multiple groups are given, each
        source variable is looked up across all groups in order and the first
        match is used.
    dst_var:
        Destination zarr array to write into.
    dst_name:
        Output variable name (used only for error messages).
    """
    groups: list[zarr.Group] = [src_groups] if isinstance(src_groups, zarr.Group) else src_groups

    if isinstance(entry, str):
        src_var = _find_var(entry, groups, dst_name)
        dst_var[region] = src_var[region]
    else:
        src_names = resolve_sources(entry)
        arrays: list[np.ndarray] = [
            _find_var(src_name, groups, dst_name)[region] for src_name in src_names
        ]
        dst_var[region] = entry(*arrays)
