#!/usr/bin/python3 -u

import argparse
import codecs
import fileinput
import os
import shutil
import sys
from pathlib import Path

def print_cr(string):
    print(string + '\r')

def is_comment(line):
    return line.lstrip().startswith(';')

def patch_file(file, encoding, append_after, new_string):
    with codecs.open(file, encoding=encoding) as ofile:
        if append_after not in ofile.read():
            ofile.close()
            with codecs.open(file, 'a', encoding=encoding) as ofile:
                ofile.write('\r\n')
                ofile.write(append_after)
                ofile.write('\r\n')
        ofile.close()
    line_injected = False
    with fileinput.input(file, encoding=encoding, inplace=True) as ofile:
        in_section = False
        found_line = False
        inject_line = False
        for line in ofile:
            if not in_section and line.strip() == append_after:
                in_section = True
            if in_section:
                if line == new_string:
                    # the line is already there
                    found_line = True
                if line.strip() != append_after and line.startswith('[') and not found_line:
                    # this is the next section and we did not find the line
                    inject_line = True
            if inject_line:
                print_cr(new_string)
                print_cr(line.rstrip())
                in_section = False
                inject_line = False
                line_injected = True
            else:
                # stripping the comments avoids that the file becomes too big for setup to load (>64k)
                if not is_comment(line):
                    print_cr(line.rstrip())
    if not line_injected:
        with codecs.open(file, 'a', encoding=encoding) as ofile:
            ofile.write(new_string)
            ofile.write('\r\n')

sections_to_copy = ['[display]', '[oemdisks]', '[pointing.device]']

def parse_oemsetup_file(file):
    oemsetup_data = []
    copy_disks_section = False
    with codecs.open(file, encoding='cp437') as ofile:
        if '[oemdisks]' not in ofile.read():
            copy_disks_section = True
        ofile.seek(0)
        section = False;
        for line in ofile.readlines():
            if is_comment(line):
                continue
            if line.lstrip().startswith('['):
                section = line.strip();
                if section == '[disks]':
                    if copy_disks_section:
                        section = '[oemdisks]'
                    else:
                        section = False
            else:
                if section in sections_to_copy and line.strip() != '':
                    oemsetup_data.append((section, line.strip()))
                if section == '[display]':
                    optional_work_section = line.split(',')[-1].strip()
                    if optional_work_section:
                        sections_to_copy.append('['+optional_work_section+']')
                else:
                    continue
        return oemsetup_data

parser = argparse.ArgumentParser(description="""Script to patch drivers into Windows 3.1 installation files.
The source direction should contain a Windows 3.1 driver with accompanying oemsetup.inf.
The destination directory the complete set of Windows 3.1 installation files.
There is a special option to add a DOSEMU2 machine type.
Repeating this script should not have averse effects (idempotent).""")
parser.add_argument("-s", "--source",      help="Path to driver directory", type=str)
parser.add_argument("-d", "--destination", help="Path to Windows 3.1 installation files", type=str)
parser.add_argument("-c", "--codepage",    help="Codepage used by Windows 3.1 installation files", type=str)

args = parser.parse_args()

if not args.source or not args.destination:
    print("Missing arguments, see -h")
    sys.exit(1)

driverinfo_file = os.path.join(args.source, 'oemsetup.inf')
if not Path(driverinfo_file).exists():
    print('Missing oemsetup.inf, source needs to point to win31 driver.')
    sys.exit(1)

setupinfo_file = os.path.join(args.destination, 'setup.inf')
if not Path(setupinfo_file).exists():
    print('Missing setup.inf, destination needs to point to win31 install.')
    sys.exit(1)

driver_data = parse_oemsetup_file(driverinfo_file)
for entry in driver_data:
    patch_file(setupinfo_file, args.codepage, entry[0], entry[1])

EXTENSION_LIST = [ '386', '3gr', 'drv', 'exe' ]
for extension in EXTENSION_LIST:
    for source in Path(args.source).glob('*.' + extension):
        filename = str(source).split("/")[-1]
        if Path(os.path.join(args.destination, filename)).exists() or Path(os.path.join(args.destination, filename[:-1] + '_')).exists():
            print('Skipping already existing file: ' + filename)
        else:
            print('Copying file: ' + filename)
            shutil.copy(source, args.destination)
