"""Single, thin command-line entry point for GNN4Colliders."""

from __future__ import annotations

import logging
import os
import sys
from importlib import resources

from hydra import compose, initialize_config_dir
from omegaconf import OmegaConf

from gnn4colliders.config.application import (
    export_onnx,
    predict_or_evaluate,
    prepare,
    train,
)

_COMMANDS = {"prepare", "train", "evaluate", "predict", "export"}


def _is_main_process() -> bool:
    return int(os.environ.get("RANK", "0")) == 0


def _help(command: str | None = None) -> None:
    if command:
        print(f"Usage: gnn4colliders {command} [key=value ...]")
        print("Compose Hydra configuration with semantic overrides.")
    else:
        print(
            "Usage: gnn4colliders <prepare|train|evaluate|predict|export> "
            "[key=value ...]"
        )
        print("Use gnn4colliders <command> --help for command help.")


def _config(overrides: list[str]):
    config_resource = resources.files("gnn4colliders.configs")
    with resources.as_file(config_resource) as config_dir:
        with initialize_config_dir(version_base=None, config_dir=str(config_dir)):
            return compose(config_name="config", overrides=overrides)


def main(argv: list[str] | None = None) -> int:
    args = list(sys.argv[1:] if argv is None else argv)
    if not args or args[0] in {"-h", "--help"}:
        _help()
        return 0
    command = args.pop(0)
    if command not in _COMMANDS:
        print(
            f"unknown command {command!r}; choose from {', '.join(sorted(_COMMANDS))}",
            file=sys.stderr,
        )
        return 2
    if "--help" in args or "-h" in args:
        _help(command)
        return 0
    try:
        config = _config(args)
        logging.basicConfig(
            level=getattr(logging, str(config.logging.level).upper(), logging.INFO)
        )
        if command == "prepare":
            print(f"prepared cache: {prepare(config)}")
        elif command == "train":
            if _is_main_process():
                print(f"saved checkpoint: {train(config)}")
        elif command == "export":
            if _is_main_process():
                print(f"wrote ONNX model: {export_onnx(config)}")
        elif command == "evaluate":
            result = predict_or_evaluate(config, evaluate=True)
            if _is_main_process():
                print(OmegaConf.to_yaml(result.metrics))
        else:
            result = predict_or_evaluate(config)
            if _is_main_process():
                print(
                    f"wrote predictions: {config.inference.output} "
                    f"({len(result.sample_ids)} events)"
                )
        return 0
    except (ValueError, FileNotFoundError, KeyError, ImportError) as error:
        print(f"gnn4colliders: {error}", file=sys.stderr)
        return 2


__all__ = ["main"]