"""Bounded A64 integer-to-WAT compiler with validated immediate control flow.

Handles selected integer operations, NZCV arithmetic, branches, and declared
host calls. It is not a whole-program lifter. Guest registers remain 64-bit;
memory addresses are checked before narrowing to the 64 KiB prototype memory.
Straight-line exports use run(ctx); CFG exports use run(ctx, instruction_budget).
Budget exhaustion returns status 3 with partial register/memory effects; it is
an abort, not a resumable checkpoint. Incoming NZCV is not part of this ABI.
"""
import struct


class UnsupportedInstruction(ValueError):
    pass


def signed(value, bits):
    return value - (1 << bits) if value & (1 << (bits - 1)) else value


def const(value):
    return f"(i64.const {signed(value & ((1 << 64) - 1), 64)})"


class Compiler:
    def __init__(self, hosts=None):
        self.hosts = hosts or {}
        self.instructions = []
        self.calls = []

    def reg(self, index, width=64, sp=False):
        value = f"(local.get ${'sp' if index == 31 else 'x' + str(index)})" if index != 31 or sp else const(0)
        return value if width == 64 else f"(i64.and {value} {const(0xffffffff)})"

    def put(self, index, value, width=64, sp=False):
        if index == 31 and not sp:
            return f"(drop {value})"
        if width == 32:
            value = f"(i64.and {value} {const(0xffffffff)})"
        return f"(local.set ${'sp' if index == 31 else 'x' + str(index)} {value})"

    def address(self, base, offset):
        return f"(i64.add {self.reg(base, sp=True)} {const(offset)})"

    def load(self, size, address):
        op = {1: 'i64.load8_u', 2: 'i64.load16_u', 4: 'i64.load32_u', 8: 'i64.load'}[size]
        return f"({op} (call $addr {address} {const(size)}))"

    def store(self, size, address, value):
        op = {1: 'i64.store8', 2: 'i64.store16', 4: 'i64.store32', 8: 'i64.store'}[size]
        return f"({op} (call $addr {address} {const(size)}) {value})"

    def arithmetic(self, rd, left, right, width, subtract, flags, sp=False):
        operation = 'sub' if subtract else 'add'
        mask = const((1 << width) - 1)
        result = [f'(local.set $a {left})', f'(local.set $b {right})',
                  f'(local.set $result (i64.and (i64.{operation} (local.get $a) (local.get $b)) {mask}))']
        if flags:
            a, b, r = '(local.get $a)', '(local.get $b)', '(local.get $result)'
            overflow = f'(i64.and (i64.xor {a} {b}) (i64.xor {a} {r}))' if subtract else \
                f'(i64.and (i64.xor {a} {r}) (i64.xor {b} {r}))'
            result += [f'(local.set $n (i32.wrap_i64 (i64.shr_u {r} {const(width-1)})))',
                       f'(local.set $z (i64.eqz {r}))',
                       f'(local.set $c (i64.{"ge_u" if subtract else "lt_u"} {a if subtract else r} {b if subtract else a}))',
                       f'(local.set $v (i32.wrap_i64 (i64.shr_u {overflow} {const(width-1)})))']
        result.append(self.put(rd, '(local.get $result)', width, sp=sp and not flags))
        return result

    def condition(self, cond):
        n, z, c, v = ['(local.get $' + name + ')' for name in ('n', 'z', 'c', 'v')]
        conditions = [z, c, n, v, f'(i32.and {c} (i32.eqz {z}))',
                      f'(i32.eq {n} {v})', f'(i32.and (i32.eq {n} {v}) (i32.eqz {z}))', '(i32.const 1)']
        result = conditions[cond >> 1]
        return f'(i32.eqz {result})' if cond & 1 and cond < 14 else result

    def decode(self, word, pc):
        rd, rn = word & 31, (word >> 5) & 31
        width = 64 if word >> 31 else 32
        if word == 0xd65f03c0:
            return ["(br $done)"], 'ret'
        # Narrow 128-bit load/store copy idiom; compile() validates the adjacent
        # matching pair and forbids control-flow entry into its store half.
        if word & 0xffe00c00 == 0x3cc00000:
            address = self.address(rn, signed((word >> 12) & 511, 9))
            return [f'(local.set $base {address})',
                    f'(local.set $qlo {self.load(8, "(local.get $base)")})',
                    f'(local.set $qhi {self.load(8, "(i64.add (local.get $base) (i64.const 8))")})'], 'copy128-load'
        if word & 0xffc00000 == 0x3d800000:
            address = self.address(rn, ((word >> 10) & 4095) * 16)
            return [f'(local.set $base {address})',
                    self.store(8, '(local.get $base)', '(local.get $qlo)'),
                    self.store(8, '(i64.add (local.get $base) (i64.const 8))', '(local.get $qhi)')], 'copy128-store'
        # Conditional select family (CSEL/CSINC/CSINV/CSNEG), including CSET.
        if word & 0x3fe00800 == 0x1a800000:
            rm, cond = (word >> 16) & 31, (word >> 12) & 15
            alternative = self.reg(rm, width)
            invert, increment = bool(word & (1 << 30)), bool(word & (1 << 10))
            if invert:
                alternative = f'(i64.xor {alternative} {const((1 << width)-1)})'
            if increment:
                alternative = f'(i64.add {alternative} (i64.const 1))'
            return [self.put(rd, f'(select {self.reg(rn, width)} {alternative} {self.condition(cond)})', width)], 'select-flags'
        if word & 0xfc000000 == 0x94000000:
            target = pc + signed(word & 0x3ffffff, 26) * 4
            if target not in self.hosts:
                raise UnsupportedInstruction(f"undeclared call to {target:#x}")
            host = self.hosts[target]
            self.calls.append({'site': hex(pc), 'target': hex(target), 'host': host})
            return [self.put(30, const(pc + 4)),
                    f"(local.set $status (call ${host} {self.reg(0)} {self.reg(1)} {self.reg(2)}))",
                    "(br_if $done (local.get $status))"], 'bl-host'
        # ADD/SUB immediate, including CMP/CMN aliases.
        if word & 0x1f000000 == 0x11000000:
            if word & (1 << 23):
                raise UnsupportedInstruction('reserved ADD/SUB immediate')
            immediate = ((word >> 10) & 4095) << (12 if word & (1 << 22) else 0)
            flags = bool(word & (1 << 29))
            return self.arithmetic(rd, self.reg(rn, width, sp=True), const(immediate), width,
                                   bool(word & (1 << 30)), flags, sp=True), 'arithmetic' + ('-flags' if flags else '')
        if word & 0x1f200000 == 0x0b000000:
            shift, amount = (word >> 22) & 3, (word >> 10) & 63
            if shift == 3 or amount >= width:
                raise UnsupportedInstruction('reserved arithmetic shift')
            source = self.reg((word >> 16) & 31, width)
            if shift == 2 and width == 32:
                source = f'(i64.extend_i32_u (i32.shr_s (i32.wrap_i64 {source}) (i32.const {amount})))'
            elif amount:
                source = f'(i64.{["shl", "shr_u", "shr_s"][shift]} {source} {const(amount)})'
            source = f'(i64.and {source} {const((1 << width)-1)})'
            flags = bool(word & (1 << 29))
            return self.arithmetic(rd, self.reg(rn, width), source, width,
                                   bool(word & (1 << 30)), flags), 'arithmetic' + ('-flags' if flags else '')
        # Move-wide immediate (MOVZ/MOVN/MOVK).
        if word & 0x1f800000 == 0x12800000:
            operation, shift = (word >> 29) & 3, ((word >> 21) & 3) * 16
            if operation == 1 or shift >= width:
                raise UnsupportedInstruction('reserved move-wide encoding')
            value = ((word >> 5) & 65535) << shift
            if operation == 0:
                expression = const((~value) & ((1 << width) - 1))
            elif operation == 2:
                expression = const(value)
            else:
                mask = (~(65535 << shift)) & ((1 << width) - 1)
                expression = f"(i64.or (i64.and {self.reg(rd, width)} {const(mask)}) {const(value)})"
            return [self.put(rd, expression, width)], 'move-wide'
        # Logical shifted register, including the MOV register alias.
        if word & 0x1f000000 == 0x0a000000:
            op, shift_kind = (word >> 29) & 3, (word >> 22) & 3
            amount, rm = (word >> 10) & 63, (word >> 16) & 31
            if op == 3 or shift_kind == 2 or amount >= width:
                raise UnsupportedInstruction('flag-setting, arithmetic shift, or reserved logical encoding')
            source = self.reg(rm, width)
            if amount:
                if shift_kind == 3 and width == 32:
                    source = f"(i64.extend_i32_u (i32.rotr (i32.wrap_i64 {source}) (i32.const {amount})))"
                else:
                    shift = {0: 'shl', 1: 'shr_u', 3: 'rotr'}[shift_kind]
                    source = f"(i64.{shift} {source} {const(amount)})"
            if word & (1 << 21):
                source = f"(i64.xor {source} {const((1 << width) - 1)})"
            operation = {0: 'and', 1: 'or', 2: 'xor'}[op]
            return [self.put(rd, f"(i64.{operation} {self.reg(rn, width)} {source})", width)], 'logical-register'
        # Unsigned bitfield. Covers UBFX, UBFIZ and immediate logical shifts.
        if word & 0x7f800000 == 0x53000000:
            if ((word >> 22) & 1) != (width == 64):
                raise UnsupportedInstruction('reserved UBFM width')
            r, s = (word >> 16) & 63, (word >> 10) & 63
            if max(r, s) >= width:
                raise UnsupportedInstruction('reserved UBFM range')
            value = self.reg(rn, width)
            if s >= r:
                value = f"(i64.and (i64.shr_u {value} {const(r)}) {const((1 << (s - r + 1)) - 1)})"
            else:
                value = f"(i64.shl (i64.and {value} {const((1 << (s + 1)) - 1)}) {const(width - r)})"
            return [self.put(rd, value, width)], 'unsigned-bitfield'
        # Integer load/store with unsigned scaled immediate.
        if word & 0x3b000000 == 0x39000000 and not word & (1 << 26):
            size, op = 1 << (word >> 30), (word >> 22) & 3
            if op not in (0, 1):
                raise UnsupportedInstruction('signed/prefetch load variant')
            address = self.address(rn, ((word >> 10) & 4095) * size)
            code = self.put(rd, self.load(size, address)) if op else self.store(size, address, self.reg(rd))
            return [code], ('load' if op else 'store') + str(size * 8)
        # Integer unscaled or pre/post-indexed immediate accesses.
        if word & 0x3b200000 == 0x38000000 and not word & (1 << 26):
            size, op, mode = 1 << (word >> 30), (word >> 22) & 3, (word >> 10) & 3
            if op not in (0, 1) or mode == 2:
                raise UnsupportedInstruction('signed/prefetch/unprivileged unscaled access')
            if mode in (1, 3) and rn != 31 and rn == rd:
                raise UnsupportedInstruction('unscaled access with unpredictable writeback overlap')
            offset = signed((word >> 12) & 511, 9)
            result = []
            if mode == 3:
                result.append(self.put(rn, self.address(rn, offset), sp=True))
            result.append(f'(local.set $base {self.address(rn, offset if mode == 0 else 0)})')
            address = '(local.get $base)'
            result.append(self.put(rd, self.load(size, address)) if op else self.store(size, address, self.reg(rd)))
            if mode == 1:
                result.append(self.put(rn, self.address(rn, offset), sp=True))
            return result, 'unscaled-access'
        # Integer pairs, including SP pre/post-indexed stack saves.
        if word & 0x3a000000 == 0x28000000 and not word & (1 << 26):
            opc, mode, rt2 = word >> 30, (word >> 23) & 3, (word >> 10) & 31
            if opc not in (0, 2) or mode == 0:
                raise UnsupportedInstruction('non-temporal/signed/reserved pair')
            if mode in (1, 3) and rn != 31 and rn in (rd, rt2):
                raise UnsupportedInstruction('pair with constrained-unpredictable writeback overlap')
            size = 8 if opc == 2 else 4
            offset = signed((word >> 15) & 127, 7) * size
            result = []
            if mode == 3:
                result.append(self.put(rn, self.address(rn, offset), sp=True))
            effective_offset = offset if mode == 2 else 0
            # Snapshot the base so register destinations cannot alter later addresses.
            result.append(f"(local.set $base {self.address(rn, effective_offset)})")
            for index, register in enumerate((rd, rt2)):
                address = f"(i64.add (local.get $base) {const(index * size)})"
                if word & (1 << 22):
                    result.append(self.put(register, self.load(size, address)))
                else:
                    result.append(self.store(size, address, self.reg(register)))
            if mode == 1:
                result.append(self.put(rn, self.address(rn, offset), sp=True))
            return result, ('load' if word & (1 << 22) else 'store') + '-pair'
        raise UnsupportedInstruction(f"unsupported opcode 0x{word:08x}")

    def compile(self, body, address):
        self.instructions = []
        self.calls = []
        if not body or len(body) % 4:
            raise ValueError('A64 input must contain complete instructions')
        words = list(struct.unpack('<' + 'I' * (len(body) // 4), body))
        for i, word in enumerate(words):
            if word & 0xffe00c00 == 0x3cc00000:
                if i + 1 >= len(words) or words[i+1] & 0xffc0001f != 0x3d800000 | (word & 31):
                    raise UnsupportedInstruction('128-bit load requires adjacent matching store')
            if word & 0xffc00000 == 0x3d800000:
                if not i or words[i-1] & 0xffe00c1f != 0x3cc00000 | (word & 31):
                    raise UnsupportedInstruction('128-bit store requires adjacent matching load')
                for j, candidate in enumerate(words):
                    branch = self.branch(candidate, address+j*4)
                    if branch and branch[0] == address+i*4:
                        raise UnsupportedInstruction('control flow enters 128-bit copy halfway')
        if any(self.branch(word, address + i * 4) is not None for i, word in enumerate(words)):
            return self.compile_cfg(words, address)
        code = []
        saw_ret = False
        defined_flags = False
        for index, (word,) in enumerate(struct.iter_unpack('<I', body)):
            pc = address + index * 4
            if saw_ret:
                raise UnsupportedInstruction('bytes after terminal RET require a control-flow model')
            try:
                emitted, kind = self.decode(word, pc)
            except UnsupportedInstruction as error:
                raise UnsupportedInstruction(f'{pc:#x}: {error}') from error
            if kind == 'select-flags' and ((word >> 12) & 15) < 14 and not defined_flags:
                raise UnsupportedInstruction(f'conditional select reads undefined NZCV at {pc:#x}')
            if kind == 'arithmetic-flags':
                defined_flags = True
            elif kind == 'bl-host':
                defined_flags = False
            self.instructions.append({'address': hex(pc), 'word': f'{word:08x}', 'kind': kind})
            code += [f';; {pc:#x}: {word:08x} {kind}', *emitted]
            saw_ret = kind == 'ret'
        if not saw_ret:
            raise UnsupportedInstruction('missing terminal RET')
        registers = [f'x{i}' for i in range(31)] + ['sp']
        imports = '\n'.join(f'(import "env" "{name}" (func ${name} (param i64 i64 i64) (result i32)))'
                            for name in sorted(set(self.hosts.values())))
        declarations = '\n'.join(f'(local ${name} i64)' for name in registers)
        loads = '\n'.join(f'(local.set ${name} (i64.load offset={i * 8} (local.get $ctx)))'
                          for i, name in enumerate(registers))
        stores = '\n'.join(f'(i64.store offset={i * 8} (local.get $ctx) (local.get ${name}))'
                           for i, name in enumerate(registers))
        return f'''(module
{imports}
(memory (export "memory") 1 1)
(func $addr (param $value i64) (param $size i64) (result i32)
  (if (i64.gt_u (local.get $value) (i64.sub (i64.const 65536) (local.get $size)))
    (then unreachable))
  (i32.wrap_i64 (local.get $value)))
(func (export "run") (param $ctx i32) (result i32)
{declarations}
(local $base i64) (local $status i32)
(local $a i64) (local $b i64) (local $result i64) (local $qlo i64) (local $qhi i64)
(local $n i32) (local $z i32) (local $c i32) (local $v i32)
{loads}
(block $done
{chr(10).join(code)}
)
{stores}
(local.get $status)))
'''

    def branch(self, word, pc):
        """Return (target, condition) for supported immediate branches."""
        if word & 0xfc000000 == 0x14000000:
            return pc + signed(word & 0x3ffffff, 26) * 4, None
        if word & 0xff000010 == 0x54000000:
            cond = word & 15
            return pc + signed((word >> 5) & 0x7ffff, 19) * 4, self.condition(cond)
        if word & 0x7e000000 == 0x34000000:
            width = 64 if word >> 31 else 32
            zero = f'(i64.eqz {self.reg(word & 31, width)})'
            condition = f'(i32.eqz {zero})' if word & (1 << 24) else zero
            return pc + signed((word >> 5) & 0x7ffff, 19) * 4, condition
        if word & 0x7e000000 == 0x36000000:
            bit = ((word >> 31) << 5) | ((word >> 19) & 31)
            zero = f'(i64.eqz (i64.and {self.reg(word & 31)} {const(1 << bit)}))'
            condition = f'(i32.eqz {zero})' if word & (1 << 24) else zero
            return pc + signed((word >> 5) & 0x3fff, 14) * 4, condition
        return None

    def compile_cfg(self, words, address):
        """Validated instruction-block dispatcher; guest loops consume a caller budget.

        Each instruction is currently a block. This deliberately simple CFG is
        correct before block coalescing/structured-loop optimization is added.
        """
        end = address + len(words) * 4
        nodes = []
        reachable = set()
        successors = {}
        pending = [address]
        while pending:
            pc = pending.pop()
            if pc in reachable:
                continue
            if not address <= pc < end or (pc - address) % 4:
                raise UnsupportedInstruction(f'control flow outside declared body: {pc:#x}')
            reachable.add(pc)
            word = words[(pc - address) // 4]
            branch = self.branch(word, pc)
            if branch:
                edges = [branch[0]]
                if branch[1] is not None:
                    edges.append(pc + 4)
            elif word != 0xd65f03c0:
                edges = [pc + 4]
            else:
                edges = []
            successors[pc] = edges
            pending.extend(edges)
        if len(reachable) != len(words):
            raise UnsupportedInstruction('unreachable bytes require explicit code/data classification')
        if not any(word == 0xd65f03c0 for word in words):
            raise UnsupportedInstruction('missing RET in control-flow body')
        for i, word in enumerate(words):
            pc = address + i * 4
            branch = self.branch(word, pc)
            if branch:
                target, condition = branch
                next_pc = f'(i32.const {(target-address)//4})'
                if condition is not None:
                    next_pc = f'(select {next_pc} (i32.const {i+1}) {condition})'
                emitted = [f'(local.set $pc {next_pc})', '(br $dispatch)']
                kind = 'branch-immediate'
            else:
                emitted, kind = self.decode(word, pc)
                emitted += [f'(local.set $pc (i32.const {i+1}))', '(br $dispatch)']
            self.instructions.append({'address': hex(pc), 'word': f'{word:08x}', 'kind': kind})
            nodes.append(f'(if (i32.eq (local.get $pc) (i32.const {i})) (then\n' + '\n'.join(emitted) + '\n))')
        # Require flags to be defined on every path before a conditional read.
        # Host calls invalidate flags: their ARM64 ABI does not preserve NZCV.
        incoming = {pc: pc != address for pc in reachable}
        outgoing = dict(incoming)
        changed = True
        while changed:
            changed = False
            for i, instruction in enumerate(self.instructions):
                pc = address + i * 4
                predecessors = [p for p, edges in successors.items() if pc in edges]
                value = pc != address and all(outgoing[p] for p in predecessors)
                after = True if instruction['kind'] == 'arithmetic-flags' else \
                    False if instruction['kind'] == 'bl-host' else value
                if (incoming[pc], outgoing[pc]) != (value, after):
                    incoming[pc], outgoing[pc] = value, after
                    changed = True
        for i, word in enumerate(words):
            cond = (word & 15) if word & 0xff000010 == 0x54000000 else \
                ((word >> 12) & 15) if self.instructions[i]['kind'] == 'select-flags' else 14
            if cond < 14 and not incoming[address+i*4]:
                raise UnsupportedInstruction(f'conditional operation reads undefined NZCV at {address+i*4:#x}')
        registers = [f'x{i}' for i in range(31)] + ['sp']
        imports = '\n'.join(f'(import "env" "{name}" (func ${name} (param i64 i64 i64) (result i32)))'
                            for name in sorted(set(self.hosts.values())))
        declarations = '\n'.join(f'(local ${name} i64)' for name in registers)
        loads = '\n'.join(f'(local.set ${name} (i64.load offset={i * 8} (local.get $ctx)))'
                          for i, name in enumerate(registers))
        stores = '\n'.join(f'(i64.store offset={i * 8} (local.get $ctx) (local.get ${name}))'
                           for i, name in enumerate(registers))
        return f'''(module
{imports}
(memory (export "memory") 1 1)
(func $addr (param $value i64) (param $size i64) (result i32)
  (if (i64.gt_u (local.get $value) (i64.sub (i64.const 65536) (local.get $size))) (then unreachable))
  (i32.wrap_i64 (local.get $value)))
(func (export "run") (param $ctx i32) (param $budget i32) (result i32)
{declarations}
(local $base i64) (local $status i32) (local $pc i32)
(local $a i64) (local $b i64) (local $result i64) (local $qlo i64) (local $qhi i64)
(local $n i32) (local $z i32) (local $c i32) (local $v i32)
{loads}
(block $done
 (loop $dispatch
  (if (i32.eqz (local.get $budget)) (then (local.set $status (i32.const 3)) (br $done)))
  (local.set $budget (i32.sub (local.get $budget) (i32.const 1)))
  {chr(10).join(nodes)}
  unreachable
 ))
{stores}
(local.get $status)))
'''