#!/usr/bin/env -S python3 -u
import argparse
import os
import seamless
import seamless.config
from seamless_config.select import get_execution, get_selected_cluster
import seamless_remote

if __name__ == "__main__":

    env_remote_clients = os.environ.get("SEAMLESS_REMOTE_CLIENTS")
    if env_remote_clients is not None:
        exit(0)

    try:
        parser = argparse.ArgumentParser(prog="seamless-init")
        parser.add_argument(
            "--stage",
            metavar="STAGE[:SUBSTAGE]",
            help="set Seamless project stage (and substage). Each project stage has independent storage",
        )
        args = parser.parse_args()
        if args.stage:
            if ":" in args.stage:
                stage, substage = args.stage.split(":", 1)
                if stage == "":
                    raise ValueError("Stage cannot be empty")
                if substage == "":
                    substage = None
            else:
                stage, substage = args.stage, None
            seamless.config.set_stage(stage, substage, workdir=os.getcwd())
        else:
            seamless.config.init(workdir=os.getcwd())
        if get_selected_cluster() is not None:
            import seamless_remote.buffer_remote
            import seamless_remote.database_remote

            seamless_remote.buffer_remote.ensure_initialized()
            seamless_remote.database_remote.ensure_initialized()

            execution = get_execution()
            if execution == "remote":
                import seamless_remote.jobserver_remote
                import seamless_remote.daskserver_remote

                seamless_remote.jobserver_remote.ensure_initialized()
                seamless_remote.daskserver_remote.ensure_initialized()

    finally:
        seamless.close()
