import os
from typing import Optional
from jinja2 import Environment, FileSystemLoader

# Setting up the jinya environment to load templates from the templates folder.
base_dir        = os.path.dirname(os.path.abspath(__file__))
template_dir    = os.path.join(base_dir, "templates")
env             = Environment(loader = FileSystemLoader(template_dir), trim_blocks = True, lstrip_blocks = True)

def generate_sbatch_script(
        nodes:          int,
        gpus_per_node:  int,
        cpus_per_task:  int,
        time_limit:     str,
        job_name:       str,
        account_name:   Optional[str],
        wandb_api_key:  Optional[str],
        hf_repo_id:     Optional[str],
        hf_dataset_id:  Optional[str],
        user_command:   str
):
    '''
    Loads the master template and renders it with user parameters.
    '''
    
    # loading the template file
    template                = env.get_template('sbatch_master.sh.j2')

    # inflating placeholders
    rendered_script         = template.render(
        nodes               = nodes,
        gpus_per_node       = gpus_per_node,
        cpus_per_task       = cpus_per_task,
        time_limit          = time_limit,
        job_name            = job_name,
        account_name        = account_name,
        wandb_api_key       = wandb_api_key,
        hf_repo_id          = hf_repo_id,
        hf_dataset_id       = hf_dataset_id,
        user_command_script = user_command
    )

    return rendered_script