#!/usr/bin/env python3

# Script to insert AmphibianDNA binaries into the build tree

import argparse
import os
import os.path as path
import time
import shutil
import re
from datetime import datetime
from subprocess import run, PIPE, DEVNULL
from machfs import Volume, Folder, File
from macresources import parse_file, make_file

########################################################################

RE1 = re.compile(rb'^(\w+)\s*=\s*(.+)')
RE2 = re.compile(rb'{(\w+)}')

def extract_makefile_defines(makefile, seed={}):
    makefile = makefile.replace(b'\r', b'\n') # tolerate all sorts of crud
    vardict = dict(seed)
    grabber = lambda m: vardict.get(m.group(1).decode('ascii'), '').encode('ascii')

    for line in text.split(b'\n'):
        m = RE1.match(line)
        if m:
            try:
                left = m.group(1).decode('ascii')
                right = RE2.sub(grabber, m.group(2)).decode('ascii')
                vardict[left] = right
            except UnicodeDecodeError:
                pass

    return vardict

########################################################################

REB = re.compile(rb'Build-date: (\d\d\d\d-\d\d-\d\d)')
def get_build_date(from_dir):
    try:
        msg = run(['git', 'rev-list', '--format=%B', '--max-count=1', 'HEAD'], cwd=from_dir, stdout=PIPE, stderr=DEVNULL, check=True)
        msg = msg.stdout
    except:
        return

    for l in msg.split(b'\n'):
        m = REB.match(l)
        if m:
            return m.group(1).decode('ascii')

def ticks_from_str(s):
    delta = datetime.strptime(datstr, '%Y-%m-%d') - datetime(1904, 1, 1)
    delta = int(delta.total_seconds())
    return delta

########################################################################
# Argparse

args = argparse.ArgumentParser(description='''
    Copy an MPW source tree into a System 7 disk image that builds on
    boot. Difficult parts of the build tree can be elegantly replaced by
    dropping pre-built objects or binaries into the AmphibianDNA folder.
    Every file in AmphibianDNA will be spliced into the build system by
    rewriting the makefile rule that seems to target it, or by direct
    copying if no rule is found. All necessary BuildResults folders are
    also created.
''')
args.add_argument('src', metavar='SOURCES', action='store', help='Source tree')
args.add_argument('-e', dest='emu', metavar='VMAC', action='store', default=None, help='path to emulator')
group = args.add_mutually_exclusive_group()
group.add_argument('--resedit', action='store_true', help='placeholder')
group.add_argument('--tomeviewer', action='store_true', help='placeholder')
group.add_argument('-c', dest='mpwcmd', metavar='CMD', action='store', default=None, help='MPW Shell command line')
args.add_argument('-v', dest='verbose', action='store_true', help='verbose')
args = args.parse_args()
treedest = path.join(args.src, 'BuildImage')
imgdest = path.join(args.src, 'BuildImage.dmg')

########################################################################

def log(*a, **kwa):
    if args.verbose:
        print(*a, **kwa)

########################################################################

log('Copying source tree', flush=True)

try:
    shutil.rmtree(treedest)
except FileNotFoundError:
    pass

myignore = shutil.ignore_patterns('BuildImage*', '.*', '*.dmg', '*.dsk', '*.sh', '*.py')

# copy2 preserves mod times, which we need to eventually allow MPW Make to work right
shutil.copytree(args.src, treedest, ignore=myignore, copy_function=shutil.copy2)

all_makefiles = []
for root, dirs, files in os.walk(treedest):
    for f in files:
        if f.lower().endswith('.make'):
            all_makefiles.append(path.join(root, f))

########################################################################

log('Creating build folders', flush=True) # I have better code for this!

main_makefiles = []
for mkfile in all_makefiles:
    with open(mkfile, 'rb') as f:
        text = f.read()

    defines = extract_makefile_defines(text)
    if 'BuildDir' not in defines: continue

    main_makefiles.append(mkfile)
    for key, macpath in defines.items():
        macpath = macpath.replace('"', '')
        if key.endswith('Dir') and macpath.startswith('BuildResults:'):
            nativepath = path.join(treedest, *macpath.split(':')[:-1])
            os.makedirs(nativepath, exist_ok=True)

########################################################################

log('Splicing amphibian DNA', flush=True) # Ugly, but keeps src tree clean

OVERDIR = 'AmphibianDNA'
try:
    overs = [n for n in os.listdir(path.join(treedest, OVERDIR)) if path.splitext(n)[1].lower() not in ('.idump', '.rdump')]
except FileNotFoundError:
    overs = []

if overs:
    remaining = list(overs)

    overs_re = '|'.join(re.escape(x) for x in overs)
    overs_re = r'^[^#\s]*\b(Thing.lib)"?\s*ƒ\s*'.replace('Thing.lib', overs_re)
    overs_re = re.compile(overs_re, re.IGNORECASE)

    for f in all_makefiles:
        mfile = open(f).read().split('\n')

        havechanged = False
        newmfile = []

        idx = -1
        while idx + 1 < len(mfile):
            idx += 1
            m = overs_re.match(mfile[idx])
            if m:
                thefile = m.group(1)
                havechanged = True
                newmfile.append('# Rule replaced at build time by ' + path.basename(__file__))
                remaining = [x for x in remaining if x.upper() != thefile.upper()]

                srcfile = '{Sources}%s:%s' % (OVERDIR, thefile)
                newmfile.append(m.group(0) + srcfile)
                newmfile.append('\tDuplicate -y {Deps} {Targ}')

                lastidx = idx # how many "old" lines should be commented out?
                if mfile[idx].endswith('∂'):
                    while lastidx + 1 < len(mfile) and mfile[lastidx + 1].endswith('∂'):
                        lastidx += 1 # capture continuations of first line
                while lastidx + 1 < len(mfile) and mfile[lastidx + 1].startswith('\t'):
                    lastidx += 1 # capture build lines starting with tab

                while idx <= lastidx:
                    newmfile.append('#\t' + mfile[idx])
                    idx += 1
            else:
                newmfile.append(mfile[idx])

        if havechanged:
            open(f, 'w').write('\n'.join(newmfile))

    if remaining: # try to find where these override files with *no build rule* should go
        found_locations = {k: [] for k in remaining}

        overs_re = '|'.join(re.escape(x) for x in remaining)
        overs_re = r'^[^#]*"({\w+}(?:\w+:)*)(Thing.lib)"'.replace('Thing.lib', overs_re)
        overs_re = re.compile(overs_re, re.IGNORECASE)

        for f in all_makefiles:
            mfile = open(f).read().split('\n')

            for line in mfile:
                m = overs_re.match(line)
                if m:
                    orig_name = next(x for x in remaining if x.upper() == m.group(2).upper())
                    found_loc = m.group(1)+m.group(2)
                    if found_loc.upper() not in (x.upper() for x in found_locations[orig_name]):
                        found_locations[orig_name].append(found_loc)

        for f in main_makefiles:
            with open(f, 'a') as fd:
                fd.write('\n# Rules created at build time by %s\n' % path.basename(__file__))
                for orig_name, found_locs in found_locations.items():
                    if len(found_locs) == 1:
                        remaining = [x for x in remaining if x != orig_name]
                        fd.write(found_locs[0])
                        fd.write(' ƒ {Sources}%s:%s\n' % (OVERDIR, orig_name))
                        fd.write('\tDuplicate -y {Deps} {Targ}\n')

    diag = 'Successfully spliced: %d/%d' % (len(overs)-len(remaining), len(overs))
    if remaining: diag += '; Failed: ' + ' '.join(remaining)
    log(diag)

########################################################################

log('Loading source tree into memory', flush=True)

vol = Volume()
vol.name = 'Disk'
vol['Src'] = Folder()
vol['Src'].read_folder(treedest, date=0x90000000, mpw_dates=True)

########################################################################

log('Inserting System and apps', flush=True)

def folder(name):
    return path.join(path.dirname(path.abspath(__file__)), name)

def pstring(string):
    string = string.encode('mac_roman')
    return bytes([len(string)]) + string

# We need a System Folder
vol['System Folder'] = Folder()
vol['System Folder'].read_folder(folder('System-7.0.1'), date=0x80000000)

if 1: # has MPW
    vol['MPW'] = Folder()
    vol['MPW'].read_folder(folder('MPW-3.2.3'), date=0x80000000)

    # Clear Worksheet window and reset its position and size
    vol['MPW']['Worksheet'].data = vol['MPW']['Worksheet'].rsrc = b''

if 1: # has apps in Apple Menu
    vol['System Folder']['Apple Menu Items'] = Folder()
    vol['System Folder']['Apple Menu Items'].read_folder(folder('Apps'), date=0x80000000)

# Tell MPW what to do when it starts (can use {Src} in their -c command)
bootscript = 'Set Src Disk:Src:; Export Src; Directory {Src}\r'
if args.mpwcmd: bootscript += 'Set Exit 0; %s; Quit\r' % args.mpwcmd
vol['MPW']['UserStartup'].data = bootscript.encode('mac_roman')

# Evil tuning: more RAM for MPW because some tools do not use temp mem
mpw = list(parse_file(vol['MPW']['MPW Shell'].rsrc))
for resource in mpw:
    if resource.type == b'SIZE':
        resource.data = resource.data[:2] + (0x300000).to_bytes(4, byteorder='big') * 2 + resource.data[10:]
vol['MPW']['MPW Shell'].rsrc = make_file(mpw)

# Patch the Process Manager to shut down when the last visible app quits,
# which it normally does when the startup app isn't a 'FNDR'.
# At that point, replace ShutDwnUserChoice() with ShutDwnPower().
# Also, follow a special path to the startup app, instead of Finder name global.
if args.mpwcmd is not None:
    startapp = 'Disk:MPW:MPW Shell'
elif args.resedit:
    startapp = 'Disk:System Folder:Apple Menu Items:ResEdit'
elif args.tomeviewer:
    startapp = 'Disk:System Folder:Apple Menu Items:TomeViewer'
else:
    startapp = None

sys = list(parse_file(vol['System Folder']['System'].rsrc))
for resource in sys:
    if resource.type != b'scod': continue # "System CODE" resources = Process Mgr

    resource.data = resource.data.replace(b'FNDR', b'LUSR')
    resource.data = resource.data.replace(b'\x3F\x3C\x00\x05\xA8\x95', b'\x3F\x3C\x00\x01\xA8\x95')

    if startapp: # Change arg 3 of FSMakeFSSpec(BootVol, 0, FinderName, &mySpec) to custom path
        data = bytearray(resource.data)
        pea = b'\x48\x78\x02\xE0' # Find 'PEA $2E0' (accesses the Str15 FinderName)...
        str_offset = None
        while pea in data:
            if not str_offset:
                while len(data) % 2: data.append(0)
                str_offset = len(data)
                data.extend(pstring(startapp))

            offset = data.index(pea)
            data[offset:offset+2] = b'\x48\x7A' # ...replace with PC-relative 'PEA startapp'
            data[offset+2:offset+4] = (str_offset - (offset+2)).to_bytes(2, byteorder='big')
        resource.data = bytes(data)
vol['System Folder']['System'].rsrc = make_file(sys)

# Patch the Finder to give it a Quit menu item
finder = list(parse_file(vol['System Folder']['Finder'].rsrc))
for resource in finder:
    if resource.type == b'fmnu' and b'Put Away' in resource.data:
        data = bytearray(resource.data)
        data[3] += 2
        data.extend(b'xxx0\0\0\0\0\x01\x2D')
        data.extend(b'quit\x81\x00\x51\x00\x04Quit\0')
        resource.data = bytes(data)
vol['System Folder']['Finder'].rsrc = make_file(finder)

########################################################################

log('Writing out disk image')
with open(imgdest, 'wb') as f:
    f.write(vol.write(256*1024*1024, align=4096, desktopdb=False))

if not args.emu: exit()
log('Starting emulator...', flush=True)

run([args.emu, imgdest], check=True)

########################################################################

log('Slurping output', flush=True)

vol = Volume()
with open(imgdest, 'rb') as f:
    vol.read(f.read())
vol['Src'].write_folder(args.src)

if args.mpwcmd is not None:
    wsheet = vol['MPW']['Worksheet'].data.decode('mac_roman').replace('\r', '\n')
    print(wsheet, end='')