"""Shared parts of the differential tests: Unicorn reference, operand generators, build and run of harness.c.

Copyright (c) 2026 Peiikai (MIT; see THIRD_PARTY_NOTICES.md).
"""
import os
import shutil
import struct
import subprocess
import sys
from pathlib import Path

import capstone
import unicorn as uc
from unicorn import arm64_const as A

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / 'recomp'))

BUF, BUF_SIZE = 0x7200000000, 0x40000
UC_INVALID = (uc.UC_ERR_INSN_INVALID, uc.UC_ERR_EXCEPTION)

# Record layouts of struct In / struct Out in harness.c.
IN_FMT = '<QQ32QQQ'
IN_TAIL = '<QQ'
OUT_SIZE = 32 * 8 + 4 * 8 + 32 * 16

MSYS_BIN = Path(r'C:\msys64\ucrt64\bin')

F32 = [0, 0x80000000, 0x7f800000, 0xff800000, 0x7fc00000, 0xffc00000, 0x7fa00001, 0xff812345, 0x7fd00abc,
       0x00000001, 0x80000001, 0x007fffff, 0x807fffff, 0x00400000, 0x00800000, 0x80800000, 0x00800001, 0x7f7fffff,
       0xff7fffff, 0x3f800000, 0xbf800000, 0x3f000000, 0xbf000000, 0x40000000, 0x40400000, 0x3fc00000, 0xbfc00000,
       0x40200000, 0x40100000, 0x4f000000, 0xcf000000, 0x4f800000, 0x5f000000, 0xdf000000, 0x5f800000, 0x4effffff,
       0x4f7fffff, 0x5effffff, 0x5f7fffff, 0x3effffff, 0x3f7fffff, 0x7effffff, 0x01000000, 0x3e800000, 0x7e800000,
       0x7e7fffff, 0x00200000, 0x001fffff, 0x00200001, 0x47800000, 0x477fe000, 0x477ff000, 0x38800000, 0x33800000,
       0x33000000, 0x387fc000, 0x3f7ff000, 0x46fffe00, 0x80000000 | 0x00200000, 0x7f000000, 0x7f7ffffe]
F64 = [0, 1 << 63, 0x7ff0000000000000, 0xfff0000000000000, 0x7ff8000000000000, 0xfff8000000000000,
       0x7ff4000000000001, 0xfff0000000012345, 0x7ffc00000abcdef0, 1, (1 << 63) | 1, 0x000fffffffffffff,
       0x0008000000000000, 0x0010000000000000, 0x8010000000000000, 0x7fefffffffffffff, 0xffefffffffffffff,
       0x3ff0000000000000, 0xbff0000000000000, 0x3fe0000000000000, 0x4000000000000000, 0x4008000000000000,
       0x3ff8000000000000, 0xbff8000000000000, 0x4004000000000000, 0x41e0000000000000, 0xc1e0000000000000,
       0x41f0000000000000, 0x43e0000000000000, 0xc3e0000000000000, 0x43f0000000000000, 0x41dfffffffc00000,
       0x41efffffffe00000, 0x43dfffffffffffff, 0x43efffffffffffff, 0x3fdfffffffffffff, 0x7fdfffffffffffff,
       0x7fd0000000000000, 0x0004000000000000, 0x0003ffffffffffff, 0x47efffffe0000000, 0x47efffff00000000,
       0x36a0000000000000, 0x3810000000000000, 0x380fffffffffffff, 0x47f0000000000000, 0x3e70000000000000,
       0x3f0ffc0000000000, 0x40effc0000000000, 0x40effe0000000000, 0x3e6ffffffff00000, 0x369fffffffffffff,
       0x3810000010000000, 0x381000000fffffff, 0x3e80000000000000]
F16 = [0, 0x8000, 0x7c00, 0xfc00, 0x7e00, 0xfe00, 0x7d00, 0x7e55, 0x0001, 0x8001, 0x03ff, 0x0400, 0x7bff, 0xfbff,
       0x3c00, 0xbc00, 0x3800, 0x4000, 0x3e00, 0x7fff, 0xffff]

_md = capstone.Cs(capstone.CS_ARCH_ARM64, capstone.CS_MODE_ARM)


def dis(w, pc=0):
    d = next(_md.disasm(struct.pack('<I', w), pc), None)
    return f'{d.mnemonic} {d.op_str}' if d else '?'


def mnemonic(w):
    return dis(w).split(' ')[0]


def rand_lane(r, sz):
    """Random interesting integer lane of 8 << sz bits."""
    bits_ = 8 << sz
    k = r.random()
    if k < 0.35:
        return r.choice([0, 1, 2, 3, (1 << bits_) - 1, (1 << bits_) - 2, 1 << (bits_ - 1), (1 << (bits_ - 1)) - 1,
                         (1 << (bits_ - 1)) + 1, 0x80 % (1 << bits_), 0x7f, bits_ - 1, bits_, bits_ + 1,
                         (1 << bits_) - bits_, 1 << (bits_ // 2)])
    if k < 0.5:
        return r.getrandbits(r.randint(1, bits_))
    if k < 0.6:
        return ((1 << bits_) - r.getrandbits(r.randint(1, bits_ - 1))) & ((1 << bits_) - 1)
    return r.getrandbits(bits_)


def rand_f32(r):
    k = r.random()
    if k < 0.4:
        return r.choice(F32)
    if k < 0.8:
        e = r.choice([r.randint(110, 145), r.randint(0, 255), r.randint(120, 135)])
        return (r.getrandbits(1) << 31) | (e << 23) | (r.getrandbits(23) if r.random() < 0.7 else
                                                        r.getrandbits(4) << 19)
    return struct.unpack('<I', struct.pack('<f', r.choice([r.uniform(-4, 4), float(r.randint(-300, 300)) + 0.5,
                                                           float(r.randint(-5, 5)), r.uniform(-1e6, 1e6)])))[0]


def rand_f64(r):
    k = r.random()
    if k < 0.4:
        return r.choice(F64)
    if k < 0.8:
        e = r.choice([r.randint(1000, 1050), r.randint(0, 2047), r.randint(1020, 1030), r.randint(0, 40),
                      r.randint(2000, 2047), r.randint(860, 920)])
        return (r.getrandbits(1) << 63) | (e << 52) | (r.getrandbits(52) if r.random() < 0.7 else
                                                        r.getrandbits(6) << 46)
    return struct.unpack('<Q', struct.pack('<d', r.choice([r.uniform(-4, 4), float(r.randint(-300, 300)) + 0.5,
                                                           r.uniform(-1e12, 1e12)])))[0]


def rand_v(r):
    k = r.random()
    out = 0
    if k < 0.3:
        for i in range(4):
            out |= rand_f32(r) << (32 * i)
    elif k < 0.55:
        for i in range(2):
            out |= rand_f64(r) << (64 * i)
    elif k < 0.6:
        for i in range(8):
            out |= (r.choice(F16) if r.random() < 0.5 else r.getrandbits(16)) << (16 * i)
    elif k < 0.9:
        sz = r.randrange(4)
        for i in range(16 >> sz):
            out |= rand_lane(r, sz) << ((8 << sz) * i)
    else:
        out = r.getrandbits(128)
    if r.random() < 0.15:  # upper half zero, like most 64-bit results
        out &= (1 << 64) - 1
    return out


def rand_x(r):
    k = r.random()
    if k < 0.3:
        return r.choice([0, 1, 2, 0xFFFFFFFFFFFFFFFF, 0x7FFFFFFFFFFFFFFF, 0x8000000000000000, 0xFFFFFFFF,
                         0x7FFFFFFF, 0x80000000, 0x100000000, 0xFFFFFFFF80000000, 0x8000000000000001, 0xFFFFFF,
                         0x1000001, 0x20000000000001, 0xFFFFFFFFFF000001, 63, 64, 31, 32, 0xFF, 0xFFFF])
    if k < 0.45:
        return rand_f32(r) | (r.getrandbits(32) << 32 if r.random() < 0.3 else 0)
    if k < 0.6:
        return rand_f64(r)
    if k < 0.75:
        return r.getrandbits(r.randint(1, 64))
    if k < 0.8:
        return (-r.getrandbits(r.randint(1, 63))) & 0xFFFFFFFFFFFFFFFF
    return r.getrandbits(64)


def rand_fpcr(r):
    k = r.random()
    if k < 0.55:
        return 0
    v = 0
    if r.random() < 0.4:
        v |= 1 << 25  # DN
    if r.random() < 0.4:
        v |= 1 << 24  # FZ
    if r.random() < 0.4:
        v |= r.randrange(4) << 22  # RMode
    if r.random() < 0.15:
        v |= 1 << 26  # AHP
    return v


class Emu:
    """Unicorn as a Cortex-A57, the Switch's CPU, unless another model is given (LSE atomics need
    UC_CPU_ARM64_MAX). Images are (base, bytes) pairs mapped at their guest address."""

    def __init__(self, images, model=A.UC_CPU_ARM64_A57):
        mu = uc.Uc(uc.UC_ARCH_ARM64, uc.UC_MODE_ARM)
        mu.ctl_set_cpu_model(model)
        mu.mem_map(BUF, BUF_SIZE)
        for base, data in images:
            mu.mem_map(base, (len(data) + 0xFFFF) & ~0xFFFF)
            mu.mem_write(base, bytes(data))
        mu.reg_write(A.UC_ARM64_REG_CPACR_EL1, 0x300000)  # enable FP/SIMD
        self.mu = mu
        self.xregs = [A.UC_ARM64_REG_X0 + i for i in range(29)] + [A.UC_ARM64_REG_X29, A.UC_ARM64_REG_X30]
        self.qregs = [A.UC_ARM64_REG_Q0 + i for i in range(32)]
        self.out_regs = self.xregs + [A.UC_ARM64_REG_SP, A.UC_ARM64_REG_NZCV, A.UC_ARM64_REG_PC] + self.qregs

    def run(self, pc, fpcr, x, sp, nzcv, v, count=1):
        """Run count instructions from pc. Returns (x[31], sp, nzcv, pc, v[32])."""
        mu = self.mu
        mu.reg_write_batch(list(zip(self.xregs, x[:31])) + list(zip(self.qregs, v)) +
                           [(A.UC_ARM64_REG_SP, sp), (A.UC_ARM64_REG_NZCV, nzcv), (A.UC_ARM64_REG_FPCR, fpcr)])
        mu.emu_start(pc, pc + 4 * count + 4, count=count)
        o = mu.reg_read_batch(self.out_regs)
        return list(o[:31]), o[31], o[32] & 0xF0000000, o[33], list(o[34:])


def compiler():
    cc = os.environ.get('CC')
    if cc:
        return cc
    if os.name == 'nt' and (MSYS_BIN / 'clang.exe').exists():
        return str(MSYS_BIN / 'clang.exe')
    return shutil.which('clang') or 'clang'


def build_harness(work, bodies, jobs, cflags):
    """Compile one C function per test body plus harness.c into work/harness[.exe]."""
    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-unused-label', '-I', str(ROOT / 'runtime')] + cflags.split()
    n = max(1, min(jobs, len(bodies) // 200 + 1))
    per = (len(bodies) + n - 1) // n
    srcs = []
    for j in range(n):
        path = work / f'tests_{j}.c'
        with open(path, 'w') as f:
            f.write('#include "cpu.h"\n')
            for k in range(j * per, min(len(bodies), (j + 1) * per)):
                f.write(f'void t{k}(Cpu *c) {{ {bodies[k]} }}\n')
        srcs.append(path)
    with open(work / 'table.c', 'w') as f:
        f.write('#include "cpu.h"\n')
        f.write(''.join(f'void t{k}(Cpu *c);\n' for k in range(len(bodies))))
        f.write('GuestFn tests[] = {' + ','.join(f't{k}' for k in range(len(bodies))) + '};\n')
        f.write(f'int ntests = {len(bodies)};\n')
    # fma.c: a --cflags=-march=x86-64-v2 run must exercise the software fma the v2 build links, not the C runtime's.
    srcs += [work / 'table.c', ROOT / 'tests' / 'diff' / 'harness.c', ROOT / 'runtime' / 'fma.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 / ('harness.exe' if os.name == 'nt' else 'harness')
    # lld when the toolchain has it beside the compiler; its default linker otherwise (an MSYS2 install without the lld
    # package stopped every differential test at "invalid linker name")
    cc_dir = os.path.dirname(shutil.which(cc, path=env.get('PATH')) or cc)
    lld = ['-fuse-ld=lld'] if any(os.path.isfile(os.path.join(cc_dir, n)) for n in ('ld.lld', 'ld.lld.exe')) else []
    subprocess.check_call([cc, *lld, '-o', str(exe)] + [str(o) for o in objs], env=env)
    return exe


def write_cases(path, cases, stubs=False):
    """cases: (k, fpcr, x[31], sp, nzcv, v[32], memchk). stubs: calls go to harness.c's call stub (rt_stub)."""
    with open(path, 'wb') as f:
        for k, fpcr, x, sp, nzcv, v, memchk in cases:
            f.write(struct.pack(IN_FMT, k, fpcr, *(list(x[:31]) + [0]), sp, nzcv))
            for q in v:
                f.write(q.to_bytes(16, 'little'))
            f.write(struct.pack(IN_TAIL, int(memchk), int(stubs)))


def run_harness(exe, work, buf0, cases, images, stubs=False):
    """Run the lifted tests. Returns (records, stderr tail or None). images: (base, path) pairs."""
    with open(work / 'buf.bin', 'wb') as f:
        f.write(buf0)
    write_cases(work / 'cases.bin', cases, stubs)
    out = work / 'out.bin'
    args = [str(exe), str(work / 'buf.bin'), str(work / 'cases.bin'), str(out)]
    for base, path in images:
        args += [str(path), f'{base:x}']
    p = subprocess.run(args, capture_output=True)
    err = None if p.returncode == 0 else f'exit {p.returncode:#x}: {p.stderr.decode(errors="replace")[-2000:]}'
    data = out.read_bytes() if out.exists() else b''
    recs = []
    for i in range(len(data) // OUT_SIZE):
        rec = data[i * OUT_SIZE:(i + 1) * OUT_SIZE]
        x = struct.unpack_from('<31Q', rec, 0)
        sp, nzcv, pc, mem = struct.unpack_from('<4Q', rec, 256)
        v = [int.from_bytes(rec[288 + 16 * j:304 + 16 * j], 'little') for j in range(32)]
        recs.append((x, sp, nzcv, pc, mem, v))
    return recs, err


def diff_state(got, want, check_pc, check_mem, check_nzcv=True):
    """got: harness record; want: (x, sp, nzcv, pc, v, mem). Returns a list of mismatches."""
    hx, hsp, hnz, hpc, hmem, hv = got
    ex, esp, enz, epc, ev, emem = want
    d = [f'x{j}: got {hx[j]:x} want {ex[j]:x}' for j in range(31) if hx[j] != ex[j]]
    if hsp != esp:
        d.append(f'sp: got {hsp:x} want {esp:x}')
    if check_nzcv and hnz != enz:
        d.append(f'nzcv: got {hnz:x} want {enz:x}')
    if check_pc and hpc != epc:
        d.append(f'pc: got {hpc:x} want {epc:x}')
    if check_mem and hmem != emem:
        d.append('memory differs')
    d += [f'v{j}: got {hv[j]:032x} want {ev[j]:032x}' for j in range(32) if hv[j] != ev[j]]
    return d