"""Command-line interface for zarr-copy."""

from __future__ import annotations

import argparse
import logging
import pathlib
import sys

import zarr
from fsspec.implementations.reference import ReferenceFileSystem
from zarr.storage import FsspecStore

from .copy import copy_group
from .transform import load_transform_spec

logger = logging.getLogger(__name__)


def _open_source_group(path: str, use_consolidated: bool = True) -> zarr.Group:
    """Open a source zarr group, with special handling for reference filesystems.

    Parameters
    ----------
    path:
        Local path, S3 URI, path to a ``.json`` / ``.parq`` reference file,
        or a ``reference::``-prefixed variant of the above.
    use_consolidated:
        Whether to use consolidated metadata. Forwarded to both regular
        stores (``zarr.open_group``) and reference filesystems (``zarr.open``).

    Returns
    -------
    zarr.Group
    """
    _PREFIX = "reference::"
    stripped = path[len(_PREFIX) :] if path.startswith(_PREFIX) else path

    if stripped.endswith(".json") or stripped.endswith(".parq"):
        storage_options: dict[str, object] = {}
        if stripped.endswith(".parq"):
            storage_options = {"remote_protocol": "file", "lazy": True}
        ref_fs = ReferenceFileSystem(stripped, asynchronous=True, storage_options=storage_options)
        store = FsspecStore(fs=ref_fs, path="")
        ref_kwargs: dict[str, object] = {}
        if not use_consolidated:
            ref_kwargs["use_consolidated"] = False
        return zarr.open(store=store, mode="r", zarr_format=2, **ref_kwargs)

    open_kwargs: dict[str, object] = {}
    if not use_consolidated:
        open_kwargs["use_consolidated"] = False
    return zarr.open_group(path, mode="r", **open_kwargs)


def _split_dimname_chunksize(s: str) -> tuple[str, int]:
    """Parse a 'dimname=chunksize' string into a (str, int) pair."""
    parts = s.split("=", 1)
    if len(parts) != 2:
        raise argparse.ArgumentTypeError(f"Expected 'dimname=chunksize', got {s!r}")
    dimname, chunksize = parts
    try:
        return str(dimname), int(chunksize)
    except ValueError:
        raise argparse.ArgumentTypeError(f"Chunk size must be an integer, got {chunksize!r}")


def build_parser() -> argparse.ArgumentParser:
    """Build and return the argument parser."""
    parser = argparse.ArgumentParser(
        prog="zarr-copy",
        description=(
            "Copy a zarr dataset while optionally rechunking variables. "
            "Both local paths and S3 URIs (s3://bucket/path) are supported. "
            "When --transform is given, multiple source datasets can be specified; "
            "the last positional argument is always the destination."
        ),
    )
    parser.add_argument(
        "zarr_paths",
        type=str,
        nargs="+",
        metavar="PATH",
        help=(
            "One or more source paths followed by the destination path. "
            "The last argument is always the destination. "
            "Multiple sources are only allowed with --transform."
        ),
    )
    parser.add_argument(
        "-v",
        dest="variables",
        type=str,
        nargs="+",
        default=None,
        metavar="VAR",
        help="Variable names to copy.  Copies all variables when omitted.",
    )
    parser.add_argument(
        "--rechunk",
        type=_split_dimname_chunksize,
        nargs="+",
        default=None,
        metavar="DIM=SIZE",
        help=(
            "Rechunk one or more dimensions, e.g. --rechunk time=1 lat=256. "
            "Requires _ARRAY_DIMENSIONS metadata on each variable "
            "(standard for xarray/CF zarr stores)."
        ),
    )
    parser.add_argument(
        "--double-remap",
        type=str,
        default=None,
        metavar="DIM",
        help=(
            "Use the two-pass copy strategy for the given dimension. "
            "Reduces peak memory usage when rechunking involves large "
            "dimension reorderings."
        ),
    )
    parser.add_argument(
        "--tmp-path",
        type=pathlib.Path,
        default=None,
        metavar="DIR",
        help=(
            "Directory for temporary files used by --double-remap. "
            "Defaults to the system temp directory."
        ),
    )
    parser.add_argument(
        "--transform",
        type=pathlib.Path,
        default=None,
        metavar="FILE",
        help=(
            "Path to a Python config file that defines a module-level 'transform' dict "
            "mapping output variable names to source variable names (str) or callables. "
            "When provided, --variables is ignored and output variables are determined "
            "entirely by the transform spec. Required when more than one source is given."
        ),
    )
    parser.add_argument(
        "--no-consolidated",
        action="store_true",
        default=False,
        dest="no_consolidated",
        help=(
            "Open source datasets with use_consolidated=False. "
            "Use this when the source store has no consolidated metadata "
            "or when the consolidated metadata is stale."
        ),
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        default=False,
        help="Enable verbose logging (DEBUG level).",
    )
    return parser


def main(argv: list[str] | None = None) -> int:
    """Entry point for the zarr-copy command."""
    parser = build_parser()
    args = parser.parse_args(argv)

    log_level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(
        level=log_level,
        format="%(asctime)s %(levelname)-8s %(name)s: %(message)s",
        datefmt="%H:%M:%S",
    )

    if len(args.zarr_paths) < 2:
        print(
            "zarr-copy: error: At least one source and one destination path are required.",
            file=sys.stderr,
        )
        return 1

    *src_paths, zarr_out = args.zarr_paths

    # Multiple sources require --transform
    if len(src_paths) > 1 and args.transform is None:
        print("zarr-copy: error: Multiple source paths require --transform.", file=sys.stderr)
        return 1

    logger.info("Source(s): %s", src_paths)
    logger.info("Destination: %s", zarr_out)

    # Open source groups
    src_groups: list[zarr.Group] = []
    for src_path in src_paths:
        try:
            src_groups.append(
                _open_source_group(src_path, use_consolidated=not args.no_consolidated)
            )
        except Exception as exc:
            print(f"Error opening source zarr '{src_path}': {exc}", file=sys.stderr)
            return 1

    # Open destination group
    try:
        dst_group = zarr.open_group(zarr_out, mode="a")
    except Exception as exc:
        print(f"Error opening destination zarr: {exc}", file=sys.stderr)
        return 1

    # Parse rechunk dict
    rechunk = dict(args.rechunk) if args.rechunk else None

    # Load optional transform spec
    transform = None
    if args.transform is not None:
        try:
            transform = load_transform_spec(args.transform)
        except (FileNotFoundError, ValueError) as exc:
            print(f"Error loading transform spec: {exc}", file=sys.stderr)
            return 1
        logger.info("Loaded transform spec: %d output variable(s)", len(transform))

    try:
        copy_group(
            src_groups[0],
            dst_group,
            variables=args.variables,
            rechunk=rechunk,
            double_remap_dim=args.double_remap,
            tmp_path=args.tmp_path,
            transform=transform,
            extra_src_groups=src_groups[1:] if len(src_groups) > 1 else None,
        )
    except KeyError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1
    except Exception as exc:
        logger.exception("Unexpected error")
        print(f"Unexpected error: {exc}", file=sys.stderr)
        return 1

    return 0


if __name__ == "__main__":
    sys.exit(main())
