"""Installed command-line interface for FlowPK workflows."""

from __future__ import annotations

import argparse
import sys
from importlib import import_module
from pathlib import Path
from typing import Sequence

from pff_pk.training.runtime import bootstrap

_COMMANDS = {
    "run": "pff_pk.pipeline.experiment",
    "train": "pff_pk.pipeline.train",
    "posttrain": "pff_pk.pipeline.curriculum",
    "consolidate": "pff_pk.pipeline.consolidate",
}

_EVALUATION_COMMANDS = {
    "annotate-lenuzza-censoring": "pff_pk.commands.annotate_lenuzza_censoring",
    "sample-generation": "pff_pk.commands.generate_empirical_samples",
    "plot-trajectories": "pff_pk.commands.plot_empirical_trajectories_from_samples",
    "plot-vpc": "pff_pk.commands.plot_empirical_vpc_from_samples",
    "plot-tsne": "pff_pk.commands.plot_empirical_tsne",
    "evaluate-generation": "pff_pk.commands.compare_generation_models",
    "evaluate-validation": "pff_pk.commands.evaluate_validation_bank",
    "evaluate-vpc": "pff_pk.commands.evaluate_pharmpy_vpc",
    "upload-vpc": "pff_pk.commands.upload_vpc_artifacts",
    "upload-empirical": "pff_pk.commands.upload_empirical_artifacts",
}


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="pff-pk",
        description="Train, sample, evaluate, plot, and publish FlowPK runs.",
    )
    parser.add_argument(
        "command",
        choices=(*_COMMANDS, *_EVALUATION_COMMANDS),
        help="workflow to run",
    )
    parser.add_argument("arguments", nargs=argparse.REMAINDER)
    return parser


def main(argv: Sequence[str] | None = None) -> None:
    args = build_parser().parse_args(argv)
    bootstrap(Path.cwd())
    if args.command in _COMMANDS:
        module = import_module(_COMMANDS[args.command])
        module.main(args.arguments)
        return
    module = import_module(_EVALUATION_COMMANDS[args.command])
    previous = sys.argv
    try:
        sys.argv = [f"pff-pk {args.command}", *args.arguments]
        module.main()
    finally:
        sys.argv = previous


if __name__ == "__main__":
    main()