#!/usr/bin/env python3
"""Submit one SLURM job per variable to rechunk a zarr dataset.

Usage
-----
    uvx --from zarr-copy zarr-copy-slurm <zarr_in> <zarr_out> [options]

The script inspects the source zarr group, collects the list of variables
to process, and submits one ``sbatch`` job per variable.  Each job calls
``zarr-copy`` with ``-v <varname>`` so variables are rechunked independently
and in parallel.

When ``--transform`` is supplied the output variables are taken from the
transform spec rather than the source group.  A single-entry spec file is
written to the log directory for each output variable, and each job is given
``--transform <file>`` instead of ``-v <varname>``.

Examples
--------
Rechunk all variables, splitting time into single steps:

    uvx --from zarr-copy zarr-copy-slurm /data/input.zarr /data/output.zarr --rechunk time=1

Same, but only for selected variables:

    uvx --from zarr-copy zarr-copy-slurm /data/input.zarr /data/output.zarr \\
        -v temperature salinity \\
        --rechunk time=1 lat=256 lon=256

Use the double-remap strategy and store temporaries on a scratch filesystem:

    uvx --from zarr-copy zarr-copy-slurm /data/input.zarr /data/output.zarr \\
        --rechunk time=1 \\
        --double-remap time \\
        --tmp-path /scratch/$USER

Rename and compute output variables via a transform spec:

    uvx --from zarr-copy zarr-copy-slurm /data/era5.zarr /data/cmip.zarr \\
        --transform transform.py \\
        --rechunk time=1 lat=256 lon=256
"""

from __future__ import annotations

import argparse
import inspect
import logging
import subprocess
import sys
import textwrap
from pathlib import Path
from typing import Callable

import zarr

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

logger = logging.getLogger(__name__)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    # --- zarr-copy passthrough arguments ------------------------------------
    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 require --transform."
        ),
    )
    parser.add_argument(
        "-v",
        dest="variables",
        type=str,
        nargs="+",
        default=None,
        metavar="VAR",
        help=(
            "Variables to rechunk. Submits a job for every variable when omitted. "
            "Ignored when --transform is provided."
        ),
    )
    parser.add_argument(
        "--rechunk",
        type=str,
        nargs="+",
        default=None,
        metavar="DIM=SIZE",
        help="Rechunk specification forwarded to zarr-copy, e.g. --rechunk time=1 lat=256.",  # noqa: E501
    )
    parser.add_argument(
        "--double-remap",
        type=str,
        default=None,
        metavar="DIM",
        help="Forwarded to zarr-copy --double-remap.",
    )
    parser.add_argument(
        "--tmp-path",
        type=str,
        default=None,
        metavar="DIR",
        help="Forwarded to zarr-copy --tmp-path.",
    )
    parser.add_argument(
        "--transform",
        type=Path,
        default=None,
        metavar="FILE",
        help=(
            "Python config file defining output variables and how to compute them. "
            "One job is submitted per output variable in the spec. "
            "Forwarded to zarr-copy --transform."
        ),
    )

    parser.add_argument(
        "--no-consolidated",
        action="store_true",
        default=False,
        dest="no_consolidated",
        help=(
            "Open source datasets with use_consolidated=False. "
            "Forwarded to zarr-copy in each job script."
        ),
    )

    # --- SLURM resource arguments -------------------------------------------
    parser.add_argument(
        "--cpus",
        type=int,
        default=4,
        metavar="N",
        help="CPUs per job (default: 4).",
    )
    parser.add_argument(
        "--mem",
        type=str,
        default="32G",
        metavar="MEM",
        help="Memory per job (default: 32G).",
    )
    parser.add_argument(
        "--walltime",
        type=str,
        default="04:00:00",
        metavar="HH:MM:SS",
        help="Wall-clock time limit per job (default: 04:00:00).",
    )
    parser.add_argument(
        "--log-dir",
        type=Path,
        default=Path("logs"),
        metavar="DIR",
        help="Directory for SLURM stdout/stderr logs (default: ./logs).",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print the sbatch scripts without submitting them.",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        default=False,
        help="Pass --verbose to each zarr-copy job so tracebacks appear in the error log.",
    )

    return parser


def _entry_to_source(entry: str | Callable) -> str:
    """Serialise a transform spec entry to a Python source string.

    Strings are emitted as quoted literals; callables are reconstructed from
    their source code via :func:`inspect.getsource` so they survive being
    written to a new file.  If source inspection fails (e.g. for built-in
    functions), a ``ValueError`` is raised.
    """
    if isinstance(entry, str):
        return repr(entry)
    try:
        src = inspect.getsource(entry)
    except (OSError, TypeError) as exc:
        raise ValueError(
            f"Cannot serialise callable {entry!r} to a transform spec file: {exc}"
        ) from exc

    # getsource may return the full line a lambda appears on (e.g. when it is
    # defined inline inside a dict literal).  Extract just the lambda
    # expression by finding the first "lambda" keyword and then scanning
    # forward for the matching end of the expression.
    src = textwrap.dedent(src).strip()

    lambda_idx = src.find("lambda")
    if lambda_idx == -1:
        # Not a lambda — return as-is (regular function definition).
        return src

    # Walk forward from the lambda, tracking bracket depth.  The expression
    # ends when depth returns to 0 and we hit a comma, closing bracket, or
    # end of string.
    depth = 0
    end = len(src)
    for i, ch in enumerate(src[lambda_idx:], start=lambda_idx):
        if ch in "([{":
            depth += 1
        elif ch in ")]}":
            if depth == 0:
                end = i
                break
            depth -= 1
        elif ch == "," and depth == 0:
            end = i
            break

    return src[lambda_idx:end].rstrip()


def write_single_var_transform_spec(
    dst_name: str,
    entry: str | Callable,
    dest_dir: Path,
) -> Path:
    """Write a single-entry transform spec Python file and return its path.

    Parameters
    ----------
    dst_name:
        Output variable name (the dict key).
    entry:
        String or callable from the original transform spec.
    dest_dir:
        Directory in which to write the file.

    Returns
    -------
    Path
        Path to the written ``.py`` file.
    """
    dest_dir.mkdir(parents=True, exist_ok=True)
    path = dest_dir / f"transform_{dst_name}.py"
    entry_src = _entry_to_source(entry)
    path.write_text(
        f"# Auto-generated single-variable transform spec for '{dst_name}'\n"
        f"transform = {{\n"
        f"    {dst_name!r}: {entry_src},\n"
        f"}}\n"
    )
    return path


def build_job_script(
    varname: str,
    src_paths: list[str],
    zarr_out: str,
    rechunk: list[str] | None,
    double_remap: str | None,
    tmp_path: str | None,
    cpus: int,
    mem: str,
    walltime: str,
    log_dir: Path,
    transform_file: Path | None = None,
    no_consolidated: bool = False,
    verbose: bool = False,
) -> str:
    """Return an sbatch script as a string for a single variable."""

    zarr_copy_cmd = [sys.executable, "-m", "zarr_copy"] + src_paths + [zarr_out]

    if transform_file is not None:
        zarr_copy_cmd += ["--transform", str(transform_file)]
    else:
        zarr_copy_cmd += ["-v", varname]

    if rechunk:
        zarr_copy_cmd += ["--rechunk"] + rechunk
    if double_remap:
        zarr_copy_cmd += ["--double-remap", double_remap]
    if tmp_path:
        zarr_copy_cmd += ["--tmp-path", tmp_path]
    if no_consolidated:
        zarr_copy_cmd += ["--no-consolidated"]
    if verbose:
        zarr_copy_cmd += ["--verbose"]

    cmd = " ".join(zarr_copy_cmd)
    log_prefix = log_dir / f"rechunk_{varname}"

    return f"""\
#!/bin/bash
#SBATCH --job-name=rechunk_{varname}
#SBATCH --cpus-per-task={cpus}
#SBATCH --mem={mem}
#SBATCH --time={walltime}
#SBATCH --output={log_prefix}_%j.out
#SBATCH --error={log_prefix}_%j.err

set -euo pipefail

echo "[${{SLURM_JOB_ID}}] Starting: {varname}"
{cmd}
echo "[${{SLURM_JOB_ID}}] Done: {varname}"
"""


def submit_job(script: str, dry_run: bool) -> str | None:
    """Submit *script* via sbatch and return the job ID, or None on dry-run.

    Returns None if sbatch fails, after printing the error message to stderr.
    """
    if dry_run:
        print(script)
        print("-" * 60)
        return None

    try:
        result = subprocess.run(
            ["sbatch", "--parsable"],
            input=script,
            capture_output=True,
            text=True,
            check=True,
        )
    except subprocess.CalledProcessError as exc:
        msg = exc.stderr.strip() or exc.stdout.strip()
        print(f"sbatch failed (exit {exc.returncode}): {msg}", file=sys.stderr)
        return None
    return result.stdout.strip()


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    if len(args.zarr_paths) < 2:
        print(
            "zarr-copy-slurm: 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-slurm: error: Multiple source paths require --transform.", file=sys.stderr)
        return 1

    # --- Determine variable list and optional per-var transform files -------
    transform_files: dict[str, Path] = {}

    if args.transform is not None:
        # Load the full spec to enumerate output variable names.
        from .transform import load_transform_spec

        try:
            spec = load_transform_spec(args.transform)
        except (FileNotFoundError, ValueError) as exc:
            print(f"Error loading transform spec: {exc}", file=sys.stderr)
            return 1

        variables = list(spec.keys())
        if not variables:
            print("Transform spec is empty — nothing to submit.", file=sys.stderr)
            return 1

        # Write single-entry spec files now (before job submission) so that
        # the files exist on disk when the SLURM jobs start.
        spec_dir = args.log_dir / "transform_specs"
        for dst_name, entry in spec.items():
            try:
                transform_files[dst_name] = write_single_var_transform_spec(
                    dst_name, entry, spec_dir
                )
                logger.debug("Wrote per-variable transform spec: %s", transform_files[dst_name])
            except ValueError as exc:
                print(f"Error serialising transform entry '{dst_name}': {exc}", file=sys.stderr)
                return 1

    else:
        # Original behaviour: derive variable list from source group.
        try:
            open_kwargs: dict[str, object] = {}
            if args.no_consolidated:
                open_kwargs["use_consolidated"] = False
            src = zarr.open_group(src_paths[0], mode="r", **open_kwargs)
        except Exception as exc:
            print(f"Error opening source zarr: {exc}", file=sys.stderr)
            return 1

        variables = args.variables or list(src.array_keys())
        if not variables:
            print("No variables found in source group.", file=sys.stderr)
            return 1

        # Validate requested variable names against the source.
        for varname in list(variables):
            if varname not in src:
                print(
                    f"Warning: '{varname}' not found in source group, skipping.",
                    file=sys.stderr,
                )
                variables.remove(varname)

    # --- Create log directory and submit ------------------------------------
    if not args.dry_run:
        args.log_dir.mkdir(parents=True, exist_ok=True)

    print(f"Submitting {len(variables)} job(s) for: {', '.join(variables)}")
    logger.debug("Log directory: %s", args.log_dir)

    for varname in variables:
        script = build_job_script(
            varname=varname,
            src_paths=src_paths,
            zarr_out=zarr_out,
            rechunk=args.rechunk,
            double_remap=args.double_remap,
            tmp_path=args.tmp_path,
            cpus=args.cpus,
            mem=args.mem,
            walltime=args.walltime,
            log_dir=args.log_dir,
            transform_file=transform_files.get(varname),
            no_consolidated=args.no_consolidated,
            verbose=args.verbose,
        )

        job_id = submit_job(script, dry_run=args.dry_run)
        if job_id is None and not args.dry_run:
            return 1
        if job_id:
            logger.info("Submitted %s -> job %s", varname, job_id)
            print(f"  Submitted {varname} -> job {job_id}")

    return 0


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