import argparse
import subprocess
import sys


STAGE_COMMANDS = {
    "data": [sys.executable, "dataset_builder.py"],
    "train": [sys.executable, "train.py"],
    "evaluate": [sys.executable, "evaluate.py"],
    "all": [
        [sys.executable, "dataset_builder.py"],
        [sys.executable, "train.py"],
        [sys.executable, "evaluate.py"],
    ],
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run WorldDisasterLM pipeline stages")
    parser.add_argument("--stage", choices=["data", "train", "evaluate", "all"], default="all")
    return parser.parse_args()


def run_command(command: list[str]) -> None:
    print("Running:", " ".join(command))
    subprocess.run(command, check=True)


def main() -> None:
    args = parse_args()
    commands = STAGE_COMMANDS[args.stage]

    if args.stage == "all":
        for command in commands:
            run_command(command)
    else:
        run_command(commands)


if __name__ == "__main__":
    main()