#!/usr/bin/python3

import os
import sys
import tempfile

def main(filename):
    try:
        with open(filename, 'r', encoding='ascii') as autoboot:
            lines = [
                line.rstrip()
                for line in autoboot
            ]
    except FileNotFoundError:
        lines = []

    # Parse autoboot.txt into a mapping of section headers to mappings of
    # keys to values (all strings). This deliberately ignores blank lines or
    # any lines that are not section headers or key=value
    sections = {}
    section = '[all]'
    for line in lines:
        if '=' in line:
            key, value = line.split('=', 1)
            key = key.strip()
            value = value.strip()
            sections.setdefault(section, {})[key] = value
        elif line.startswith('['):
            section = line.rstrip()

    # If the value we want to add is already present, we're done. We shouldn't
    # edit the file unless we absolutely have to
    if sections.get('[all]', {}).get('tryboot_a_b', '') == '1':
        return 0
    if sections.get('[tryboot]', {}).get('tryboot_a_b', '') == '1':
        return 0

    # Otherwise, try and add to an existing [all] or [tryboot] section (doesn't
    # matter which). If neither exist, add an [all] section. Bear in mind that
    # autoboot.txt is limited to one sector (512 bytes) in size, so we try and
    # keep our changes as minimal as possible
    if '[all]' in sections:
        sections['[all]']['tryboot_a_b'] = '1'
    elif '[tryboot]' in sections:
        sections['[tryboot]']['tryboot_a_b'] = '1'
    else:
        sections['[all]'] = {'tryboot_a_b': '1'}

    # Re-write the file. This will necessary remove blank lines and other
    # things we ignored when parsing. But given the stringent file limit, and
    # the fact we are adding content, this is desirable
    with tempfile.NamedTemporaryFile(
        'w', dir=os.path.dirname(filename), encoding='ascii', delete=False
    ) as temp:
        for header, section in sections.items():
            temp.write(header)
            temp.write('\n')
            for key, value in section.items():
                temp.write(key)
                temp.write('=')
                temp.write(value)
                temp.write('\n')

        if temp.tell() > 512:
            os.unlink(temp.name)
            raise ValueError(f'Generated {filename} is too large')
        os.rename(temp.name, filename)


if __name__ == '__main__':
    try:
        main(*sys.argv[1:])
    except (ValueError, OSError, TypeError) as err:
        sys.stderr.write(str(err))
        sys.stderr.write('\n')
        sys.exit(1)
    else:
        sys.exit(0)
