import yaml
import os
from jinja2 import Environment, FileSystemLoader

# Configuration
SRC_DIR = 'src'
TEMPLATES_DIR = 'templates'
DIST_DIR = 'dist'
SCHEME_FILES = ['neon.yaml', 'aeon.yaml']

def load_scheme(filename):
    with open(os.path.join(SRC_DIR, filename), 'r') as f:
        return yaml.safe_load(f)

def render_template(template_path, context):
    env = Environment(loader=FileSystemLoader(TEMPLATES_DIR))
    template = env.get_template(template_path)
    return template.render(context)

def build():
    # Load all schemes first
    schemes = []
    for f in SCHEME_FILES:
        schemes.append(load_scheme(f))

    # Ensure dist directory exists
    if not os.path.exists(DIST_DIR):
        os.makedirs(DIST_DIR)

    # Walk through templates
    for root, dirs, files in os.walk(TEMPLATES_DIR):
        for file in files:
            if not file.endswith('.j2'):
                continue

            rel_path = os.path.relpath(root, TEMPLATES_DIR)
            app_name = rel_path # e.g., 'nvim', 'zed', 'kitty'
            
            template_rel_path = os.path.join(rel_path, file)
            output_filename_base = file.replace('.j2', '')

            # Create app directory in dist
            app_dist_dir = os.path.join(DIST_DIR, app_name)
            if not os.path.exists(app_dist_dir):
                os.makedirs(app_dist_dir)

            print(f"Processing template: {template_rel_path}")

            # Special case for Zed: Aggregated output (single JSON for all themes)
            if app_name == 'zed':
                context = {'schemes': schemes}
                rendered_content = render_template(template_rel_path, context)
                
                # Output file: dist/zed/apex.json (matches base name)
                output_path = os.path.join(app_dist_dir, output_filename_base)
                
                with open(output_path, 'w') as f:
                    f.write(rendered_content)
                print(f"  -> {output_path} (Aggregated)")
            
            # Default behavior: Per-scheme output
            else:
                for scheme in schemes:
                    scheme_slug = scheme['scheme'].lower().replace(' ', '-')
                    
                    # Rename output file: e.g., apex.conf -> apex-neon.conf
                    # or apex.lua -> apex-neon.lua
                    # If filename is generic 'apex', replace with slug.
                    # If filename is 'apex.lua', result 'apex-neon.lua'
                    
                    if 'apex' in output_filename_base:
                        final_filename = output_filename_base.replace('apex', f"apex-{scheme_slug.replace('apex-', '')}")
                    else:
                         final_filename = f"{scheme_slug}-{output_filename_base}"

                    # Clean up double prefixes if any (e.g. apex-apex-neon)
                    final_filename = final_filename.replace('apex-apex-', 'apex-')

                    output_path = os.path.join(app_dist_dir, final_filename)
                    
                    rendered_content = render_template(template_rel_path, scheme)
                    
                    with open(output_path, 'w') as f:
                        f.write(rendered_content)
                    print(f"  -> {output_path} ({scheme['scheme']})")

if __name__ == '__main__':
    build()