# standard library imports
import subprocess
from pathlib import Path

import git

# third-party imports
import hydra
import yaml
from omegaconf import DictConfig, OmegaConf

git_root = git.Repo(search_parent_directories=True).git.rev_parse("--show-toplevel")


@hydra.main(
    version_base=None,
    config_path=str(Path(git_root) / "configs/state_tracking"),
    config_name="state_tracking",
)
def main(cfg: DictConfig) -> None:
    output_dir = hydra.core.hydra_config.HydraConfig.get().runtime.output_dir  # type: ignore

    # set trainer output dir
    cfg.trainer.output_dir = output_dir
    cfg.trainer.project_home = git_root

    with open(Path(output_dir) / ".hydra/config.yaml", "w") as fp:
        yaml.dump(OmegaConf.to_container(cfg, resolve=True), fp)

    with open(
        Path(cfg.trainer.project_home) / "configs/state_tracking/htcondor_sweep.sh"
    ) as f:
        sweep_template = f.read()

    sweep_file = sweep_template.format(
        project_home=cfg.trainer.project_home,
        wandb_resume=cfg.logger.wandb_resume,
        wandb_project=cfg.logger.wandb_project_name,
        wandb_entity=cfg.logger.wandb_entity,
        wandb_group=cfg.logger.wandb_group,
        wandb_run_name=cfg.logger.run_name,
        config_path=str(Path(output_dir) / ".hydra"),
        wandb_dir=str(Path(output_dir)),
    )

    sweep_path = Path(output_dir) / "htcondor_sweep.sh"
    with open(sweep_path, "w") as f:
        f.write(sweep_file)

    # Make the sweep file executable
    sweep_path.chmod(0o755)

    with open(
        Path(cfg.trainer.project_home) / "configs/state_tracking/htcondor.sub"
    ) as f:
        sub_template = f.read()

    sub_file = sub_template.format(
        executable=str(Path(output_dir) / "htcondor_sweep.sh"),
        project_home=cfg.trainer.project_home,
        output_dir=cfg.trainer.output_dir,
    )

    with open(Path(output_dir) / "htcondor.sub", "w") as f:
        f.write(sub_file)

    subprocess.run(
        ["condor_submit_bid", str(500), str(Path(output_dir) / "htcondor.sub")]
    )


if __name__ == "__main__":
    main()
