import os
import yaml
from jinja2 import Environment, FileSystemLoader

def render():
    config_path = '../config/config.yaml'
    if not os.path.exists(config_path):
        print("Config file not found.")
        return

    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)

    env = Environment(loader=FileSystemLoader('templates'))
    templates = ['common.sh.j2', 'master.sh.j2', 'worker_cpu.sh.j2', 'worker_gpu.sh.j2']
    
    output_dir = "."
    os.makedirs(output_dir, exist_ok=True)

    for temp_name in templates:
        try:
            template = env.get_template(temp_name)
            output_name = temp_name.split('.')[0] + '.sh'
            rendered = template.render(config)
            
            out_path = os.path.join(output_dir, output_name)
            with open(out_path, 'w') as f:
                f.write(rendered)
            os.chmod(out_path, 0o755)
            print(f"Generated: {out_path}")
        except Exception as e:
            print(f"Error rendering {temp_name}: {e}")

if __name__ == "__main__":
    render()