"""
Typed argument dataclasses for the ``ivst`` CLI.

Each dataclass corresponds to one ``ivst`` sub-command and is parsed by
`tyro <https://brentyi.github.io/tyro/>`_.  All fields are statically typed
so IDEs and linters can validate command-line arguments before execution.
"""
import tyro
from tyro.conf import Positional
from dataclasses import dataclass, field
from typing import Union, Annotated, Literal


@dataclass
class Init:
    """Arguments for ``ivst init``.

    Generates a ``<name>.json`` DoE config file from the project's
    ``entry/__init__.py::LIST_OF_ARGS``.
    """

    #: Base name for the generated JSON file (e.g. ``doe`` → ``doe.json``).
    name: str


@dataclass
class Commit:
    """Arguments for ``ivst commit``.

    Reads ``<name>.json``, expands the parameter space, assigns case hashes,
    and generates Slurm shell scripts under ``.investigation/<commit>/shell/``.
    """

    #: Name of the DoE JSON file to commit (without ``.json``).
    name: str
    #: Git commit hash to associate with this batch of cases.
    commit: str


@dataclass
class Submit:
    """Arguments for ``ivst submit``.

    Submits pending Slurm shell scripts for a given commit.
    """

    #: Commit hash identifying which batch to submit.
    commit: str
    #: Restrict to specific batch hashes (repeatable).
    batch: Annotated[list[str], str] = field(default_factory=list)
    #: Re-submit cases that are already marked as started or done.
    force: bool = False
    #: Maximum number of jobs submitted concurrently.  ``-1`` submits all at once.
    parellels: int = -1
    #: Also submit cases in ``PENDING`` state (no shell script yet).
    include_pending: bool = False


@dataclass
class Sample:
    """Arguments for ``ivst sample``.

    Regenerates shell scripts for cases that do not yet have one.
    """

    #: Commit hash to sample from.
    commit: str
    #: Restrict to specific batch hashes (repeatable).
    batch: Annotated[list[str], str] = field(default_factory=list)
    #: Regenerate scripts even for cases already sampled.
    force: bool = False
    #: Include ``PENDING`` cases (not yet assigned a shell script).
    include_pending: bool = False


@dataclass
class Filter:
    """Arguments for ``ivst filter``.

    Inspect or mutate the set of cases for a commit.
    """

    #: Operation to perform: ``show``, ``list``, ``failed``, ``delete``, ``reset``, or ``stop``.
    action: Literal["show", "failed", "delete", "reset", "stop", "list"]
    #: Target commit hash.
    commit: str
    #: Restrict to specific batch hashes (repeatable).
    batch: Annotated[list[str], str] = field(default_factory=list)
    #: Merge cases that share the same parameters but differ only in seed.
    seed_merge: bool = False


@dataclass
class Clean:
    """Arguments for ``ivst clean``.

    Remove runtime artefacts (log files, case directories) for a commit.
    """

    #: Target commit hash.
    commit: str
    batch: Annotated[list[str], str] = field(default_factory=list)  # optional
    dry_run: bool = False


@dataclass
class Params:
    """Arguments for ``ivst params``.

    Insert additional key-value pairs into existing case configs.
    """

    #: Target commit hash.
    commit: str
    #: JSON string of parameters to insert, e.g. ``'{"lr": 0.01}'``.
    add: str
    batch: Annotated[list[str], str] = field(default_factory=list)  # optional
    dry_run: bool = False


@dataclass
class Stat:
    """Arguments for ``ivst stat``.

    Print a status summary (PENDING / SAMPLED / STARTED / DONE) for all
    cases under a commit.
    """

    #: Target commit hash.
    commit: str

    batch: Annotated[list[str], str] = field(default_factory=list)


@dataclass
class Sync:
    """Arguments for ``ivst sync``.

    Sync the current project source tree into the frozen snapshot stored under
    ``.investigation/<commit>/src/`` and into all live case directories.
    """

    #: Target commit hash.
    commit: str


@dataclass
class Analyze:
    """Arguments for ``ivst analyze``.

    Aggregate and export metric columns across cases for a commit.
    """

    #: Target commit hash.
    commit: str
    #: Metric column names to aggregate (repeatable).
    targets: Annotated[list[str], str]
    #: Output format: ``csv``, ``latex`` table, or ``print`` to stdout.
    output_format: Literal["csv", "latex", "print"]
    #: Restrict to specific batch hashes (repeatable).
    batch: Annotated[list[str], str] = field(default_factory=list)
    #: Column name to group results by (empty string means no grouping).
    output_group: str = ""
    #: Aggregation method across seeds: ``none``, ``mean``, ``min``, or ``max``.
    aggregate: str = "none"
    #: Merge cases that share the same parameters but differ only in seed.
    seed_merge: bool = False


@dataclass
class Roll:
    """Arguments for ``ivst roll``.

    Chain submitted jobs so each triggers the next on completion.
    """

    #: Target commit hash.
    commit: str
    #: Restrict to specific batch hashes (repeatable).
    batch: Annotated[list[str], str] = field(default_factory=list)


CLI = Annotated[
    Union[
        Init,
        Commit,
        Submit,
        Sample,
        Filter,
        Clean,
        Params,
        Stat,
        Sync,
        Analyze,
        Roll,
    ],
    tyro.conf.subcommand(),
]
