import os
from pathlib import Path

import htcondor

JOB_BID_SINGLE = 2000 # 100
JOB_BID_MULTI = 2000

def launch_lm_label_job(
        JOB_MEMORY,
        JOB_CPUS,
        JOB_GPUS=1,
        JOB_BID=JOB_BID_SINGLE,
        GPU_MEM=None,
        model_dir="/fast/XXXX-3/models/Qwen3-8B",
        NUM_ARTICLES=5,
):
    # Name/prefix for cluster logs related to this job
    LOG_PATH = "/fast/XXXX-3/logs/forecasting/custom_evals_vllm/"

    CLUSTER_LOGS_SAVE_DIR=Path(LOG_PATH)
    cluster_job_log_name = str(
        CLUSTER_LOGS_SAVE_DIR
        / f"$(Cluster).$(Process)"
    )

    # model_dir += "/snapshots/model/"

    # executable = 'lm_countdown.sh'
    # executable = 'launch_train_job.sh'
    executable = 'launch_train_job.sh'

    # Construct job description
    job_settings = {
        "executable": executable,
        "arguments": (
            f"{model_dir} "
            f"{NUM_ARTICLES} "
        ),
        "output": f"{cluster_job_log_name}.out",
        "error": f"{cluster_job_log_name}.err",
        "log": f"{cluster_job_log_name}.log",
        
        # "request_cpus": f"{JOB_CPUS}",  # how many CPU cores we want
        # "request_gpus": f"{JOB_GPUS}",
        # "request_memory": JOB_MEMORY,  # how much memory we want
        # "request_disk": JOB_MEMORY,
        
        "request_gpus": f"{JOB_GPUS}",
        "request_cpus": f"{max(JOB_CPUS*JOB_GPUS, 32)}",  # how many CPU cores we want
        "request_memory": f"{JOB_MEMORY*JOB_GPUS}GB",  # how much memory we want
        "request_disk": f"{JOB_MEMORY*JOB_GPUS}GB",
        
        "jobprio": f"{JOB_BID - 1000}",
        "notify_user": "XXXX-1.XXXX-2@tuebingen.mpg.de",  # otherwise one does not notice an you can miss clashes
        "notification": "error",
    }

    if GPU_MEM is not None:
        job_settings["requirements"] = f"(TARGET.CUDAGlobalMemoryMb >= {GPU_MEM}) && (CUDACapability >= 9.0)"
        # job_settings["requirements"] = f"(TARGET.CUDAGlobalMemoryMb >= {GPU_MEM}) && (CUDACapability >= 9.0) && (UtsnameNodename != 'i101') && (UtsnameNodename != 'i301') && (UtsnameNodename != 'i302') && (UtsnameNodename != 'i303') && (UtsnameNodename != 'i304') && (UtsnameNodename != 'i305') && (UtsnameNodename != 'i306') && (UtsnameNodename != 'i307') && (UtsnameNodename != 'i308')"
    else:
        job_settings["requirements"] = "CUDACapability >= 9.0"
        # job_settings["requirements"] = "(CUDACapability >= 9.0) && (UtsnameNodename != 'i101') && (UtsnameNodename != 'i301') && (UtsnameNodename != 'i302') && (UtsnameNodename != 'i303') && (UtsnameNodename != 'i304') && (UtsnameNodename != 'i305') && (UtsnameNodename != 'i306') && (UtsnameNodename != 'i307') && (UtsnameNodename != 'i308')"
        
    
    # job_settings["requirements"] = "CUDACapability >= 9.0"


    job_description = htcondor.Submit(job_settings)

    # Submit job to scheduler
    schedd = htcondor.Schedd()  # get the Python representation of the scheduler
    submit_result = schedd.submit(job_description)  # submit the job

    print(
        f"Launched experiment with cluster-ID={submit_result.cluster()}, "
        f"proc-ID={submit_result.first_proc()}")


if __name__ == "__main__":
    # from weak_models_utils import models
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_dir', type=str, default="/fast/XXXX-3/models/Qwen3-8B", help="Model directory")
    parser.add_argument('--num_articles', type=int, default=5, help="Number of articles to use per prompt")
    args = parser.parse_args()
    # parser.add_argument('--base_save_dir', type=str, required=True)  # e.g., /fast/groups/sf/ttt/evaluations/base/

    # args = parser.parse_args()
    GPU_MEM = 45000
    
    launch_lm_label_job(
        JOB_MEMORY=64,
        JOB_CPUS=8,
        JOB_GPUS=8,
        JOB_BID=JOB_BID_MULTI,
        GPU_MEM=GPU_MEM,
        model_dir=args.model_dir,
        NUM_ARTICLES=args.num_articles,
    )
        # break 
