"""Differential test of whole emitted functions: recomp/gen.py's C for a guest function against Unicorn running the
same function from its entry until it returns.

Picks leaf functions from the ExeFS modules: no calls, no branches out of the function, no SVC, BRK, system
register, exclusive or atomic instructions; with conditional branches, and loops, jump tables and SIMD where the
module has them. Emits them with the generator's own Emitter, compiles them with tests/diff/harness.c, and runs each
from random states: x0-x7 point into the scratch buffer (or hold small integers), the buffer holds pointers into
itself and random data, SP sits near the buffer's top and LR is a sentinel address. Unicorn (Cortex-A57) stops when
PC reaches the sentinel. A case is skipped when Unicorn faults, writes outside the scratch buffer (module images are
mapped read-only), leaves the function, or runs longer than --steps instructions. Compares X0-X30, SP, NZCV, V0-V31
and the scratch buffer.

Usage: func.py [--exefs DIR] [--modules main,subsdk0,sdk] [--functions N] [--cases K] [--seed S] [--jobs J]
"""
import argparse
import os
import platform
import random
import subprocess
import sys
import time
import zlib
from collections import Counter
from pathlib import Path

import numpy as np
import unicorn as uc
from unicorn import arm64_const as A

from common import BUF, BUF_SIZE, MSYS_BIN, ROOT, Emu, compiler, diff_state, rand_fpcr, rand_v, rand_x, run_harness
import gen
from layout import layout
from module import Module

SENTINEL = 0x7210000000  # LR on entry; never executed
PTR_LO, PTR_HI = BUF + 0x8000, BUF + 0x28000
STACK_TOP = BUF + 0x3F000
MIN_INSNS, MAX_INSNS = 6, 600


def forbidden_mask(w):
    """Words a picked function must not contain."""
    m = (w & 0xFC000000) == 0x94000000                    # BL
    m |= (w & 0xFFFFFC1F) == 0xD63F0000                   # BLR
    m |= (w & 0xFF000000) == 0xD4000000                   # SVC, HVC, SMC, BRK, HLT, DCPS
    sysi = (w & 0xFFC00000) == 0xD5000000                 # system: MRS, MSR, SYS (DC ZVA), hints, barriers
    hint_or_barrier = ((w & 0xFFFFF000) == 0xD5032000) | ((w & 0xFFFFF000) == 0xD5033000)
    m |= sysi & ~hint_or_barrier
    m |= (w & 0x3F000000) == 0x08000000                   # exclusives, load-acquire / store-release, CAS
    m |= (w & 0x3F200C00) == 0x38200000                   # LSE atomics
    return m


def cond_mask(w):
    return (((w & 0xFE000010) == 0x54000000) | ((w & 0x7E000000) == 0x34000000) |
            ((w & 0x7E000000) == 0x36000000))


def simd_mask(w):
    return ((w & 0x0E000000) == 0x0E000000) | (((w & 0x0A000000) == 0x08000000) & (((w >> 26) & 1) == 1))


def pick_functions(name, mod, count, r):
    """[(entry, end, tags)] of up to count leaf functions of mod, plus the Emitter that will emit them."""
    starts, _ = gen.discover(mod)
    w = mod.words
    ends = np.append(starts[1:], len(w))

    def per_fn(mask):
        cs = np.concatenate([[0], np.cumsum(mask.astype(np.int64))])
        return cs[ends] - cs[starts]

    size = ends - starts
    ok = (per_fn(forbidden_mask(w)) == 0) & (per_fn(cond_mask(w)) > 0) & (size >= MIN_INSNS) & (size <= MAX_INSNS)
    has_br = per_fn((w & 0xFFFFFC1F) == 0xD61F0000) > 0
    cand = np.flatnonzero(ok)
    with_br = [int(i) for i in cand if has_br[i]]
    rest = [int(i) for i in cand if not has_br[i]]
    r.shuffle(with_br)
    r.shuffle(rest)
    emitter = gen.Emitter(mod, set((mod.lo + starts * 4).tolist()))
    picks, tags_seen = [], Counter()
    # Functions with a br first (jump tables, fallbacks), at most a quarter of the picks.
    for pool, quota in ((with_br, count // 4), (rest, count)):
        taken = 0
        for i in pool:
            if len(picks) >= count or taken >= quota:
                break
            entry, end = mod.lo + int(starts[i]) * 4, mod.lo + int(ends[i]) * 4
            tags = check_function(emitter, entry, end)
            if tags is None:
                continue
            picks.append((entry, end, tags))
            tags_seen.update(tags)
            taken += 1
    print(f'{name}: {len(starts)} functions, {len(cand)} leaf candidates, picked {len(picks)}: '
          f'{dict(tags_seen)}', flush=True)
    return picks, emitter


def check_function(emitter, entry, end):
    """Tags of a function the test can run in isolation, or None."""
    mod = emitter.mod
    insns = [mod.lift(a) for a in range(entry, end, 4)]
    if insns[-1].kind not in gen.TERMINATORS or not any(i.kind == 'ret' for i in insns):
        return None
    tags = set()
    for k, ins in enumerate(insns):
        a = entry + 4 * k
        if ins.kind in ('bl', 'blr'):
            return None
        if ins.kind in ('b', 'bcond'):
            if not entry <= ins.target < end:
                return None
            if ins.target <= a:
                tags.add('loop')
        if ins.kind == 'br':
            jt = gen.jumptables.resolve(mod, a, entry, end)
            if jt:
                tags.add('jump_table')
            elif gen.jumptables.is_computed(mod, a, entry):
                tags.add('all_label_fallback')
            else:
                return None
    if simd_mask(mod.words[(entry - mod.lo) >> 2:(end - mod.lo) >> 2]).any():
        tags.add('simd')
    return tags


def build(work, emitted, jobs, cflags):
    """Compile the emitted functions with harness.c into work/func[.exe]; tests[k] is emitted[k]."""
    cc = compiler()
    env = dict(os.environ)
    if os.name == 'nt' and MSYS_BIN.exists():
        env['PATH'] = str(MSYS_BIN) + os.pathsep + env.get('PATH', '')
    flags = ['-O1', '-fno-strict-aliasing', '-Wno-return-type', '-Wno-unused-label', '-I', str(ROOT / 'runtime')]
    flags += cflags.split()
    groups = {}
    for name, entry, end, emitter in emitted:
        groups.setdefault(name, []).append((entry, end))
    srcs, n = [], max(1, jobs)
    for name, fns in groups.items():
        emitter = next(e for m, _, _, e in emitted if m == name)
        per = (len(fns) + n - 1) // n
        for j in range(0, len(fns), per):
            path = work / f'{name}_{j // per:03d}.c'
            path.write_text(gen.chunk_text(name, j // per, emitter, fns[j:j + per]))
            srcs.append(path)
    table = work / 'table.c'
    table.write_text('#include "cpu.h"\n' + ''.join(f'void {gen.fname(e)}(Cpu *c);\n' for _, e, _, _ in emitted) +
                     'GuestFn tests[] = {' + ','.join(gen.fname(e) for _, e, _, _ in emitted) + '};\n' +
                     f'int ntests = {len(emitted)};\n')
    srcs += [table, ROOT / 'tests' / 'diff' / 'harness.c']
    objs, procs = [], []
    for s in srcs:
        o = work / (s.stem + '.o')
        objs.append(o)
        procs.append(subprocess.Popen([cc] + flags + ['-c', str(s), '-o', str(o)], env=env))
    if any(p.wait() for p in procs):
        raise SystemExit('compile failed')
    exe = work / ('func.exe' if os.name == 'nt' else 'func')
    subprocess.check_call([cc, '-fuse-ld=lld', '-o', str(exe)] + [str(o) for o in objs], env=env)
    return exe


def make_buffer(r):
    out = bytearray()
    for _ in range(BUF_SIZE // 8):
        k = r.random()
        if k < 0.4:
            v = r.randrange(PTR_LO, PTR_HI, 8)
        elif k < 0.7:
            v = r.randrange(0, 256)
        else:
            v = r.getrandbits(64)
        out += v.to_bytes(8, 'little')
    return bytes(out)


def make_case(r):
    x = [rand_x(r) for _ in range(31)]
    for i in range(8):
        k = r.random()
        if k < 0.7:
            x[i] = r.randrange(PTR_LO, PTR_HI, 16)
        elif k < 0.9:
            x[i] = r.randrange(0, 64)
    x[30] = SENTINEL
    sp = STACK_TOP - r.randrange(0, 0x100, 16)
    return rand_fpcr(r), x, sp, r.getrandbits(4) << 28, [rand_v(r) for _ in range(32)]


class Runner:
    """Unicorn with the module images read-only and a block hook that stops when execution leaves the function."""

    def __init__(self, images, buf0):
        self.emu = Emu(images)
        mu = self.emu.mu
        for base, data in images:
            mu.mem_protect(base, (len(data) + 0xFFFF) & ~0xFFFF, uc.UC_PROT_READ | uc.UC_PROT_EXEC)
        mu.mem_map(SENTINEL, 0x1000)
        mu.mem_write(BUF, buf0)
        self.buf0 = buf0
        self.lo = self.hi = 0
        self.left = False
        mu.hook_add(uc.UC_HOOK_BLOCK, self._block)

    def _block(self, mu, address, size, user):
        if not self.lo <= address < self.hi and address != SENTINEL:
            self.left = True
            mu.emu_stop()

    def run(self, entry, end, case, steps):
        """(x, sp, nzcv, pc, v, crc of the buffer) or None if the case is skipped."""
        fpcr, x, sp, nzcv, v = case
        e, mu = self.emu, self.emu.mu
        self.lo, self.hi, self.left = entry, end, False
        mu.reg_write_batch(list(zip(e.xregs, x[:31])) + list(zip(e.qregs, v)) +
                           [(A.UC_ARM64_REG_SP, sp), (A.UC_ARM64_REG_NZCV, nzcv), (A.UC_ARM64_REG_FPCR, fpcr)])
        try:
            mu.emu_start(entry, SENTINEL, count=steps)
            ok = not self.left and mu.reg_read(A.UC_ARM64_REG_PC) == SENTINEL
        except uc.UcError:
            ok = False
        if not ok:
            mu.mem_write(BUF, self.buf0)
            return None
        o = mu.reg_read_batch(e.out_regs)
        crc = zlib.crc32(bytes(mu.mem_read(BUF, BUF_SIZE)))
        mu.mem_write(BUF, self.buf0)
        return list(o[:31]), o[31], o[32] & 0xF0000000, o[33], list(o[34:]), crc


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument('--exefs', type=Path, default=ROOT / 'build' / 'extracted' / 'v121' / 'exefs')
    ap.add_argument('--modules', default='main,subsdk0,sdk')
    ap.add_argument('--functions', type=int, default=400, help='per module')
    ap.add_argument('--cases', type=int, default=12, help='random states tried per function')
    ap.add_argument('--steps', type=int, default=200000, help='Unicorn instruction cap per case')
    ap.add_argument('--seed', type=int, default=1)
    ap.add_argument('--jobs', type=int, default=os.cpu_count() or 4)
    ap.add_argument('--show', type=int, default=20)
    ap.add_argument('--cflags', default='-march=x86-64-v3' if platform.machine().lower() in ('amd64', 'x86_64') else '')
    args = ap.parse_args()
    r = random.Random(args.seed)
    t0 = time.time()
    want = args.modules.split(',')

    work = ROOT / 'build' / 'test' / 'func'
    work.mkdir(parents=True, exist_ok=True)
    images, emitted, tags = [], [], []
    for name, base, _ in layout(args.exefs):
        if name not in want:
            continue
        mod = Module(name, args.exefs / name, base, 'c')
        img = work / f'{name}.img'
        img.write_bytes(mod.image)
        images.append((base, img, mod.image))
        picks, emitter = pick_functions(name, mod, args.functions, r)
        for entry, end, t in picks:
            emitted.append((name, entry, end, emitter))
            tags.append(t)
    print(f'picked {len(emitted)} functions ({time.time() - t0:.0f}s)', flush=True)

    t1 = time.time()
    exe = build(work, emitted, args.jobs, args.cflags)
    print(f'compiled ({time.time() - t1:.0f}s)', flush=True)

    buf0 = make_buffer(r)
    runner = Runner([(base, data) for base, _, data in images], buf0)
    cases, expect, tested, skipped = [], [], set(), 0
    t2 = time.time()
    for k, (name, entry, end, _) in enumerate(emitted):
        for _ in range(args.cases):
            case = make_case(r)
            want_state = runner.run(entry, end, case, args.steps)
            if want_state is None:
                skipped += 1
                continue
            fpcr, x, sp, nzcv, v = case
            cases.append((k, fpcr, x, sp, nzcv, v, True))
            expect.append(want_state)
            tested.add(k)
    print(f'unicorn: {len(cases)} cases on {len(tested)} functions, {skipped} skipped ({time.time() - t2:.0f}s)',
          flush=True)

    t3 = time.time()
    recs, err = run_harness(exe, work, buf0, cases, [(base, img) for base, img, _ in images])
    print(f'harness: {len(recs)} of {len(cases)} cases ({time.time() - t3:.0f}s)')
    if err:
        k = cases[len(recs)][0] if len(recs) < len(cases) else None
        where = f', stopped in {emitted[k][0]} {emitted[k][1]:#x}' if k is not None else ''
        print(f'harness failed: {err}{where}')

    fails, nfail = {}, 0
    for case, want_state, got in zip(cases, expect, recs):
        d = diff_state(got, want_state, check_pc=False, check_mem=True)
        if d:
            nfail += 1
            fails.setdefault(case[0], d[:4])
    by_tag, fail_tag = Counter(), Counter()
    for k in tested:
        by_tag.update(tags[k] or {'plain'})
        if k in fails:
            fail_tag.update(tags[k] or {'plain'})
    print(f'{len(tested)} functions, {len(cases)} cases: {nfail} failing cases, {len(fails)} failing functions '
          f'({time.time() - t0:.0f}s total)')
    print(f'  tested by tag: {dict(by_tag)}; failing by tag: {dict(fail_tag)}')
    for k, d in list(fails.items())[:args.show]:
        name, entry, end, _ = emitted[k]
        print(f'  {name:8} {gen.fname(entry)} ({(end - entry) // 4} insns, {sorted(tags[k])}): {"; ".join(d)}')
    return 1 if fails or err else 0


if __name__ == '__main__':
    sys.exit(main())