#!/usr/bin/env python3
"""setup.py - prepare the Power C improvement toolchain (the original disk is only read).

  python tools/setup.py            extract the runtime library objects into lib/ and build optlib/
  python tools/setup.py --disk     also write POWERc-opt.d81 (original + improved runtime objects)

optlib/ gets drop-in replacements (same file names as on POWERc.d81):
  cmulti.obj  faster 16x16 multiply
  c0047.obj   faster divide / remainder (same results, fewer shift steps)
  c0077.obj   the unsigned >= helper with its broken branch opcode ($80 -> $90) fixed
"""
import os, sys, subprocess, shutil, hashlib
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.normpath(os.path.join(HERE, '..'))
sys.path.insert(0, HERE)
from objfmt import Obj
from pcasm import assemble
sys.path.insert(0, os.path.join(ROOT, 'src'))
import div_patch

C1541 = r'C:\Program Files\GTK3VICE-3.10-win64\bin\c1541.exe'
ORIG = os.path.normpath(os.path.join(ROOT, '..', 'POWERc.d81'))


def run(*a):
    return subprocess.run([C1541] + list(a), capture_output=True, text=True)


def extract():
    os.makedirs(os.path.join(ROOT, 'lib'), exist_ok=True)
    out = run(ORIG, '-list').stdout
    names = []
    for ln in out.splitlines():
        if '"' in ln and ln.rstrip().endswith('seq'):
            nm = ln.split('"')[1]
            if nm.endswith(('.obj', '.o', '.l')):
                names.append(nm)
    args = []
    for nm in sorted(set(names)):
        args += ['-read', nm + ',s', os.path.join(ROOT, 'lib', nm)]
    run(ORIG, *args)
    print('extracted %d files into lib/' % len(set(names)))


def build():
    os.makedirs(os.path.join(ROOT, 'optlib'), exist_ok=True)
    cm, _ = assemble(open(os.path.join(ROOT, 'src', 'cmulti.s')).read())
    open(os.path.join(ROOT, 'optlib', 'cmulti.obj'), 'wb').write(cm.build())
    o = Obj.parse(open(os.path.join(ROOT, 'lib', 'c0047.obj'), 'rb').read())
    open(os.path.join(ROOT, 'optlib', 'c0047.obj'), 'wb').write(div_patch.patch(o).build())
    o = Obj.parse(open(os.path.join(ROOT, 'lib', 'c0077.obj'), 'rb').read())
    assert o.code[0x1d] == 0x80, 'c0077.obj is not the expected original'
    o.code[0x1d] = 0x90
    open(os.path.join(ROOT, 'optlib', 'c0077.obj'), 'wb').write(o.build())
    print('built optlib/cmulti.obj c0047.obj c0077.obj')


def disk():
    dst = os.path.join(ROOT, 'POWERc-opt.d81')
    shutil.copyfile(ORIG, dst)
    for nm in ('cmulti.obj', 'c0047.obj', 'c0077.obj'):
        run(dst, '-delete', nm)
        run(dst, '-write', os.path.join(ROOT, 'optlib', nm), nm + ',s')
    print('wrote', dst)


if __name__ == '__main__':
    h = hashlib.sha256(open(ORIG, 'rb').read()).hexdigest()
    exp = open(os.path.join(ROOT, 'POWERc.d81.sha256')).read().split()[0]
    if h != exp:
        print('WARNING: reference/POWERc.d81 differs from the recorded original hash')
    extract()
    build()
    if '--disk' in sys.argv:
        disk()