import argparse
import os
import sys

from .utils import run_command_on_all_hosts

def main() -> None:
    parser = argparse.ArgumentParser(description='A helper script to execute commands on multiple hosts of a TPU pod.')
    parser.add_argument('-i', '--include-local', action='store_true', help='include local host (127.0.0.1) in the host list')
    parser.add_argument('-c', '--clean-up', action='store_true', help='clean up temporary files generated by previous TPU processes before executing the command')
    parser.add_argument('-w', '--cwd', action='store_true', help='run the command in the current working directory, assuming the directory exists on all hosts')
    args = parser.parse_args()

    command = sys.stdin.read()
    if not command.strip():
        return

    if args.clean_up:
        command = f'sudo rm -rf /tmp/libtpu_lockfile /tmp/tpu_logs; {command}'

    if args.cwd:
        cwd = os.getcwd()
        command = f'cd {cwd}; {command}'

    run_command_on_all_hosts(command, include_local=args.include_local)
