#!/usr/bin/env python3
"""Reproduce static converter evidence and bounded Mono oracle probes (E02)."""
import argparse
import hashlib
import json
from pathlib import Path
import platform
import re
import shutil
import struct
import subprocess

from inspect_mb3dat import inspect

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / 'assets/analysis/converter'
BINARIES = {
    '2015-10-03': '66f1e3d9d53acfbe12a2be0e2a2bde64509d61d451ba7773bda68cff9d54abe5',
    '2016-03-10-c': '0789c98b4150f281f92181a3368c86f5a81e033b7ecf58b8d320c5365546b5a6',
}


def sha(data):
    return hashlib.sha256(data).hexdigest()


def binary(build):
    p = ROOT / 'assets/references/extracted' / build / 'tools/Mb3 file formats.exe'
    if sha(p.read_bytes()) != BINARIES[build]:
        raise ValueError(f'wrong converter identity: {p}')
    return p


def pe_info(data):
    pe = struct.unpack_from('<I', data, 60)[0]
    machine, sections = struct.unpack_from('<HH', data, pe + 4)
    opt = pe + 24
    opt_size = struct.unpack_from('<H', data, pe + 20)[0]
    clr = struct.unpack_from('<I', data, opt + 96 + 14 * 8)[0]
    section_table = []
    for i in range(sections):
        vsize, va, size, offset = struct.unpack_from('<4I', data, opt + opt_size + 40 * i + 8)
        section_table.append((va, max(vsize, size), offset))
    def file_offset(rva):
        for va, size, offset in section_table:
            if va <= rva < va + size:
                return offset + rva - va
        raise ValueError(f'unmapped RVA {rva:#x}')
    def cstring(offset):
        return data[offset:data.index(b'\0', offset)].decode('ascii')
    imports = []
    imp_rva = struct.unpack_from('<I', data, opt + 96 + 8)[0]
    if imp_rva:
        cursor = file_offset(imp_rva)
        while any(data[cursor:cursor + 20]):
            lookup, _, _, dll, thunk = struct.unpack_from('<5I', data, cursor)
            entries = []
            slot = file_offset(lookup or thunk)
            while True:
                entry = struct.unpack_from('<I', data, slot)[0]
                if not entry:
                    break
                entries.append({'ordinal': entry & 0xffff} if entry & 0x80000000 else cstring(file_offset(entry) + 2))
                slot += 4
            imports.append({'dll': cstring(file_offset(dll)), 'symbols': entries})
            cursor += 20
    for i in range(sections):
        vsize, va, size, offset = struct.unpack_from('<4I', data, opt + opt_size + 40 * i + 8)
        if va <= clr < va + max(vsize, size):
            c = offset + clr - va
            return {'machine': hex(machine), 'clr_rva': hex(clr), 'clr_offset': hex(c),
                    'corflags': hex(struct.unpack_from('<I', data, c + 16)[0]),
                    'entrypoint_token': hex(struct.unpack_from('<I', data, c + 20)[0]), 'pe_imports': imports}
    raise ValueError('CLR header missing')


def methods(il):
    result = []
    for match in re.finditer(r'(?ms)^\s*\.method .*?// end of method ([^\n]+)', il):
        part = match.group(0)
        rva = re.search(r'Method begins at RVA (0x[0-9a-f]+)', part)
        size = re.search(r'Code size: (\d+)', part)
        result.append({'name': match.group(1).strip(), 'rva': rva.group(1) if rva else None,
                       'il_bytes': int(size.group(1)) if size else None,
                       'il_line': il[:match.start()].count('\n') + 1})
    return result


def decompile():
    version = subprocess.check_output(['ilspycmd', '--version'], text=True).strip()
    manifest = {'tool': version, 'mono': subprocess.check_output(['mono', '--version'], text=True), 'binaries': {}}
    for build in BINARIES:
        p = binary(build)
        dest = OUT / build
        dest.mkdir(parents=True, exist_ok=True)
        for name, options in [('decompiled.cs', []), ('assembly.il', ['-il'])]:
            with (dest / name).open('wb') as stream:
                subprocess.run(['ilspycmd', '--disable-updatecheck', *options, str(p)], stdout=stream, check=True, timeout=120)
        for table in ['assembly', 'assemblyref', 'manifest', 'implmap', 'method']:
            with (dest / (table + '.txt')).open('wb') as stream:
                subprocess.run(['monodis', '--' + table, str(p)], stdout=stream, check=True, timeout=20)
        il = (dest / 'assembly.il').read_text()
        info = {'path': str(p.relative_to(ROOT)), 'sha256': sha(p.read_bytes()),
                'size': p.stat().st_size, **pe_info(p.read_bytes()), 'methods': methods(il),
                'outputs': {p.name: sha(p.read_bytes()) for p in sorted(dest.iterdir()) if p.is_file()}}
        manifest['binaries'][build] = info
    (OUT / 'static.json').write_text(json.dumps(manifest, indent=2) + '\n')
    print('decompiled and indexed both distinct converter binaries')


def i32(n):
    return struct.pack('<i', n)


def vecs(values):
    return b''.join(struct.pack('<fff', *v) for v in values)


TRIPLE = [(1, 2, 3), (4, 5, 6), (7, 8, 9)]
PLANE = 'plane(1 2 3, 4 5 6, 7 8 9)'
ANCHORS = [(u, v, u + v) for v in range(3) for u in range(3)]


def vectors_text(values):
    return ', '.join(' '.join(map(str, v)) for v in values)


def dat(version, brush_records=b'', brush_count=0, patch_records=b'', patch_count=0, computed=False, tables=b'\0\0', widths=0):
    head = b'Mb3World' + i32(version) + bytes([int(computed) | widths])
    if version == 201501280:
        return head + i32(brush_count) + i32(patch_count) + brush_records + patch_records
    return head + tables + i32(brush_count) + brush_records + i32(patch_count) + patch_records


def cases(version):
    result = []
    def add(name, src, content, dst='mb3txt', code=0, **extra):
        result.append(dict(name=name, src=src, data=content.encode('ascii') if isinstance(content, str) else content,
                           dst=dst, expected_exit=code, **extra))
    for computed in [False, True]:
        add('empty-' + str(computed).lower(), 'mb3txt', f'world(computedFeatures {str(computed).lower()})', 'mb3dat')
    text = 'world(computedFeatures false brush(' + PLANE + '))'
    add('plane-to-dat', 'mb3txt', text, 'mb3dat')
    add('plane-txt', 'mb3txt', text)
    for fmt in ['hl1map', 'hl2vmf', 'q3map']:
        add('plane-to-' + fmt, 'mb3txt', text, fmt)
    add('obj-needs-computed', 'mb3txt', text, 'wfobj', 2)
    poly = [(0, 0, 0), (1, 0, 0), (0, 1, 0)]
    computed_plane = 'world(computedFeatures true brush(plane(1 2 3, 4 5 6, 7 8 9 polygon[' + vectors_text(poly) + '])))'
    add('computed-plane', 'mb3txt', computed_plane, 'mb3dat')
    add('computed-plane-obj', 'mb3txt', computed_plane, 'wfobj')
    for computed in [False, True]:
        patch = 'patch9(3 3 2 2 anchors[' + vectors_text(ANCHORS) + ']'
        if computed:
            patch += ' points[' + vectors_text(ANCHORS) + '] normals[' + vectors_text([(0, 1, 0)] * 9) + ']'
        patch += ')'
        txt = f'world(computedFeatures {str(computed).lower()} {patch})'
        add('patch-' + str(computed).lower(), 'mb3txt', txt, 'mb3dat')
        if computed:
            add('patch-obj', 'mb3txt', txt, 'wfobj')
        else:
            add('patch-q3', 'mb3txt', txt, 'q3map')
            add('patch-hl1-dropped', 'mb3txt', txt, 'hl1map')
    base = dat(version)
    add('dat-empty', 'mb3dat', base, 'mb3dat')
    add('dat-trailing', 'mb3dat', base + b'EXTRA', 'mb3dat')
    add('wrong-version', 'mb3dat', base[:8] + i32(0) + base[12:], code=2)
    other = 201512200 if version == 201501280 else 201501280
    add('cross-version', 'mb3dat', dat(other), code=2)
    # A partial final scalar exercises the unchecked reused read buffer without an EOF string loop.
    single = dat(version, b'\0\1' + vecs(TRIPLE), 1)
    add('dat-plane', 'mb3dat', single, 'mb3dat')
    add('truncated-last-float', 'mb3dat', single[:-8] + single[-4:], 'mb3dat')
    for n in [255, 256]:
        width = 0 if n == 255 else 1
        body = bytes([width]) + struct.pack('<B' if width == 0 else '<H', n) + vecs(TRIPLE) * n
        add('planes-' + str(n), 'mb3dat', dat(version, body, 1), 'mb3dat')
    # Legacy selector=1 is signed; count 32768 immediately fails before reading plane records.
    if version == 201501280:
        add('signed-16-count', 'mb3dat', dat(version, b'\1\0\x80', 1), code=2)
    else:
        add('reserved-width-3', 'mb3dat', dat(version, tables=i32(0) + i32(0), widths=6 | 24), 'mb3dat')
        add('reserved-world-bits', 'mb3dat', base[:12] + b'\xe0' + base[13:], 'mb3dat')
        props = 'world(computedFeatures false stringProperties ["unused", "rare", "common"] uvProperties ["uv"] brush('
        props += 'plane(1 2 3, 4 5 6, 7 8 9 s 1 "r" s 2 "c" uv 0 (1 2 3 4 5, 6 7 8 9 10)) '
        props += 'plane(1 2 3, 4 5 6, 7 8 9 s 2 "")))'
        add('properties-to-dat', 'mb3txt', props, 'mb3dat')
        add('properties-to-txt', 'mb3txt', props)
        for n in [255, 256]:
            names = ','.join('"p' + str(i) + '"' for i in range(n))
            slots = ' '.join('s ' + str(i) + ' "v"' for i in range(n))
            txt = f'world(computedFeatures false stringProperties[{names}] brush(plane(1 2 3,4 5 6,7 8 9 {slots})))'
            add('properties-' + str(n), 'mb3txt', txt, 'mb3dat')
        for name, tail in [('duplicate-slot', 's 0 "a" s 0 "b"'), ('bad-slot', 's 1 "a"')]:
            add(name, 'mb3txt', 'world(computedFeatures false stringProperties["p"] brush(plane(1 2 3,4 5 6,7 8 9 ' + tail + ')))', code=2)
        add('duplicate-name', 'mb3txt', 'world(computedFeatures false stringProperties["p","p"])', code=2)
        unicode_txt = 'world(computedFeatures false stringProperties["é"] brush(plane(1 2 3,4 5 6,7 8 9 s 0 "é")))'.encode('utf-8')
        add('utf8-bom-ascii-loss', 'mb3txt', b'\xef\xbb\xbf' + unicode_txt, 'mb3dat')
        add('utf8-no-bom-mojibake', 'mb3txt', unicode_txt)
        # Duplicate binary slots are accepted with last value winning, unlike TXT.
        body = b'\0\1\0first\0\0last\0\xff' + vecs(TRIPLE)
        add('binary-duplicate-slot', 'mb3dat', dat(version, body, 1, tables=b'\1p\0\0'))
        # Reused scalar buffer turns missing patch count into the previous float's bytes (zero).
        body = b'\0\1' + vecs([(1, 2, 3), (4, 5, 6), (7, 8, 0)])
        add('truncated-patch-count', 'mb3dat', dat(version, body, 1)[:-4], 'mb3dat')
    for fmt, content in [
        ('hl1map', '{"classname" "func_detail" { (1 2 3) (4 5 6) (7 8 9) texture [1 0 0 0] [0 1 0 0] 0 1 1 }}'),
        ('hl2vmf', 'entity{"classname" "func_detail" solid{side{"plane" "(1 2 3) (4 5 6) (7 8 9)" "material" "secret"}}}'),
        ('q3map', '{"classname" "func_detail" {(1 2 3) (4 5 6) (7 8 9) texture 0 0 0 1 1 0 0 0}}')]:
        add(fmt + '-import', fmt, content, 'mb3dat')
    for fmt in ['hl1map', 'hl2vmf', 'q3map']:
        add(fmt + '-roundtrip', fmt, b'', 'mb3dat', from_case='plane-to-' + fmt)
    add('patch-dat-roundtrip', 'mb3dat', b'', 'mb3dat', from_case='patch-true')
    add('patch-q3-roundtrip', 'q3map', b'', 'mb3dat', from_case='patch-q3')
    add('inplace-dat', 'mb3dat', base, 'mb3dat', inplace=True)
    add('overwrite-dat', 'mb3txt', text, 'mb3dat', preexisting=True)
    add('computed-missing-polygon', 'mb3txt', text.replace('false', 'true'), code=2)
    add('false-ignores-polygon', 'mb3txt', computed_plane.replace('true', 'false'), 'mb3dat')
    body = b'\x04\x02' + vecs(TRIPLE) + struct.pack('<H', 256) + vecs([(0, 0, 0)] * 256)
    body += vecs(TRIPLE) + struct.pack('<H', 1) + vecs([(1, 1, 1)])
    add('contour-width-shared', 'mb3dat', dat(version, body, 1, computed=True), 'mb3dat')
    for n, selector in [(256, 1), (65536, 2)]:
        # Zero anchors isolate the dimension width without allocating huge geometry.
        body = bytes([selector << 2]) + b'\0\0' + struct.pack('<HH' if selector == 1 else '<ii', n, 2)
        add('patch-quads-' + str(n), 'mb3dat', dat(version, patch_records=body, patch_count=1), 'mb3dat')
    if version == 201512200:
        add('properties-dat-roundtrip', 'mb3dat', b'', 'mb3dat', from_case='properties-to-dat')
        add('properties-dat-to-txt', 'mb3dat', b'', from_case='properties-to-dat')
        add('utf8-txt-repeat', 'mb3txt', b'', from_case='utf8-no-bom-mojibake')
    add('bad-mode', 'MB3TXT', text, code=2)
    add('bad-syntax', 'mb3txt', 'world(', code=2)
    add('usage', 'mb3txt', text, code=1, arguments=[])
    add('missing-input', 'mb3txt', text, code=2, missing=True)
    return result


def probes():
    mono = shutil.which('mono')
    if not mono:
        raise ValueError('mono is required; enter nix develop')
    report = {'environment': {'system': platform.system(), 'release': platform.release(), 'machine': platform.machine(),
                              'mono': subprocess.check_output([mono, '--version'], text=True),
                              'cwd': '.work/converter-oracle/<build>/<case>', 'timeout_seconds': 10}, 'builds': {}}
    for build in BINARIES:
        version = 201501280 if build.startswith('2015') else 201512200
        source = binary(build)
        records = []
        for case in cases(version):
            work = ROOT / '.work/converter-oracle' / build / case['name']
            dest = OUT / 'probes' / build / case['name']
            work.mkdir(parents=True, exist_ok=True)
            dest.mkdir(parents=True, exist_ok=True)
            shutil.copyfile(source, work / 'converter.exe')
            inp, output = 'input.' + case['src'].lower(), 'output.' + case['dst']
            payload = case['data']
            if 'from_case' in case:
                previous = OUT / 'probes' / build / case['from_case']
                payload = next(previous.glob('output.*')).read_bytes()
            (work / inp).write_bytes(payload)
            if case.get('missing'):
                (work / inp).unlink()
            (work / 'conversion.log').write_bytes(b'STALE LOG\n')
            if case.get('inplace'):
                output = inp
            elif (work / output).exists():
                (work / output).unlink()
            if case.get('preexisting'):
                (work / output).write_bytes(b'PREEXISTING OUTPUT')
            args = case.get('arguments', [case['src'], inp, case['dst'], output, 'conversion.log'])
            proc = subprocess.run([mono, 'converter.exe', *args], cwd=work, capture_output=True, timeout=10)
            record = {k: v for k, v in case.items() if k != 'data'}
            record.update({'input_sha256': sha(payload), 'argv': ['mono', 'converter.exe', *args], 'binary_sha256': sha(source.read_bytes()),
                           'exit': proc.returncode, 'files': {}})
            # Fixture payload retained even in missing-file experiments.
            (dest / inp).write_bytes(payload)
            for name, data in [('stdout.txt', proc.stdout), ('stderr.txt', proc.stderr)]:
                (dest / name).write_bytes(data)
            for name in [output, 'conversion.log']:
                if (work / name).exists():
                    shutil.copyfile(work / name, dest / name)
                elif (dest / name).exists():
                    (dest / name).unlink()
            if case.get('inplace'):
                (dest / 'before.mb3dat').write_bytes(payload)
            if (work / output).exists() and case['dst'] == 'mb3dat':
                try:
                    record['inspection'] = inspect((work / output).read_bytes())
                except ValueError as exc:
                    record['inspection_error'] = str(exc)
            for p in sorted(dest.iterdir()):
                record['files'][p.name] = {'sha256': sha(p.read_bytes()), 'bytes': p.stat().st_size}
            if proc.returncode != case['expected_exit']:
                raise ValueError(f"{build}/{case['name']}: exit {proc.returncode}, expected {case['expected_exit']}; {proc.stderr!r}")
            records.append(record)
        report['builds'][build] = records
    (OUT / 'probes.json').write_text(json.dumps(report, indent=2, allow_nan=False) + '\n')
    print('captured', sum(map(len, report['builds'].values())), 'oracle probes')


def check():
    manifest = json.loads((OUT / 'static.json').read_text())
    for build, info in manifest['binaries'].items():
        binary(build)
        assert pe_info(binary(build).read_bytes())['pe_imports'] == [{'dll': 'mscoree.dll', 'symbols': ['_CorExeMain']}]
        for name, digest in info['outputs'].items():
            assert sha((OUT / build / name).read_bytes()) == digest, (build, name)
    report = json.loads((OUT / 'probes.json').read_text())
    for build, records in report['builds'].items():
        expected_cases = cases(201501280 if build.startswith('2015') else 201512200)
        assert [r['name'] for r in records] == [c['name'] for c in expected_cases]
        for case, record in zip(expected_cases, records):
            payload = case['data']
            if 'from_case' in case:
                previous = OUT / 'probes' / build / case['from_case']
                payload = next(previous.glob('output.*')).read_bytes()
            assert record['input_sha256'] == sha(payload), record['name']
            assert record['binary_sha256'] == BINARIES[build]
            assert record['exit'] == record['expected_exit'], record['name']
            if 'inspection' in record:
                output_name = record['argv'][5]
                payload_file = OUT / 'probes' / build / record['name'] / output_name
                assert inspect(payload_file.read_bytes()) == record['inspection'], record['name']
            for name, info in record['files'].items():
                p = OUT / 'probes' / build / record['name'] / name
                assert sha(p.read_bytes()) == info['sha256'] and p.stat().st_size == info['bytes'], p
        by = {r['name']: r for r in records}
        version = 201501280 if build.startswith('2015') else 201512200
        dest = OUT / 'probes' / build
        assert (dest / 'empty-false/output.mb3dat').read_bytes() == dat(version)
        assert (dest / 'plane-to-dat/output.mb3dat').read_bytes() == dat(version, b'\0\1' + vecs(TRIPLE), 1)
        assert (dest / 'dat-plane/output.mb3dat').read_bytes() == (dest / 'plane-to-dat/output.mb3dat').read_bytes()
        for name in ['hl1map-roundtrip', 'hl2vmf-roundtrip', 'q3map-roundtrip', 'overwrite-dat']:
            assert by[name]['inspection'] == by['plane-to-dat']['inspection']
        assert by['patch-dat-roundtrip']['inspection'] == by['patch-true']['inspection']
        assert by['patch-q3-roundtrip']['inspection']['patches'][0]['quads_u'] == 4
        for record in records:
            log = (OUT / 'probes' / build / record['name'] / 'conversion.log').read_bytes()
            if record['exit'] == 0 or record['name'] in ['usage', 'bad-mode']:
                assert log == b'STALE LOG\n'
            else:
                assert log.startswith(b'An exception of type ')
        assert by['dat-trailing']['inspection']['trailing_bytes'] == 0
        assert by['dat-plane']['inspection']['brushes'][0]['planes'][0]['u'] == [4, 5, 6]
        for n in [255, 256]:
            assert len(by['planes-' + str(n)]['inspection']['brushes'][0]['planes']) == n
        if build.startswith('2016'):
            d = by['properties-to-dat']['inspection']
            assert d['string_properties'] == ['common', 'rare']
            assert d['brushes'][0]['planes'][0]['uv']['0'] == [1, 2, 3, 6, 7, 8, 4, 5, 9, 10]
            for n in [255, 256]:
                assert by['properties-' + str(n)]['inspection']['string_properties'] == ['p' + str(i) for i in range(n)]
            assert by['utf8-bom-ascii-loss']['inspection']['string_properties'] == ['?']
            assert by['reserved-width-3']['inspection']['flags'] == 0
            assert by['truncated-patch-count']['exit'] == 0
            assert by['properties-dat-roundtrip']['inspection'] == by['properties-to-dat']['inspection']
            dest = OUT / 'probes' / build
            assert (dest / 'properties-to-txt/output.mb3txt').read_bytes() == (dest / 'properties-dat-to-txt/output.mb3txt').read_bytes()
            assert (dest / 'utf8-txt-repeat/output.mb3txt').read_bytes() != (dest / 'utf8-no-bom-mojibake/output.mb3txt').read_bytes()
    print('converter evidence identities, captured exits, and semantic assertions passed')


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('action', choices=['decompile', 'probes', 'check'])
    args = p.parse_args()
    try:
        {'decompile': decompile, 'probes': probes, 'check': check}[args.action]()
    except (ValueError, OSError, subprocess.SubprocessError, AssertionError) as exc:
        p.exit(1, f'{exc}\n')


if __name__ == '__main__':
    main()