"""Bounded multi-function A64 dispatcher with shared registers and guest memory.

Direct BL enters a declared guest function, B can tail-call declared code, and
RET dispatches through the guest link register. No native recursion or JIT is
used. Guest stack instructions retain their original semantics. All guest
instructions share one budget. Indirect calls and arbitrary computed returns
are unsupported. Host calls use an explicit status ABI, with an opt-in context
adapter that can return x0. Static ADRP pages must be explicitly mapped. The
zero-vector stores require a MOVI-zero definition on every incoming path.
"""
import struct
from a64_lift import Compiler, UnsupportedInstruction, const, signed


class ProgramCompiler(Compiler):
    def __init__(self, hosts=None, *, returning_hosts=(), memory_pages=1, static_pages=None):
        super().__init__(hosts)
        if type(memory_pages) is not int or not 1 <= memory_pages <= 32768:
            raise ValueError('Guest memory must contain 1..32768 fixed WASM pages')
        self.memory_pages = memory_pages
        self.returning_hosts = set(returning_hosts)
        if not self.returning_hosts <= set(self.hosts):
            raise ValueError('Returning hosts must be declared host targets')
        self.static_pages = dict(static_pages or {})
        self.zero_stores = {}
        for original, guest in self.static_pages.items():
            if original < 0 or original >= 1 << 64 or original % 4096 or guest < 0 or guest % 4096 or guest + 4096 > memory_pages*65536:
                raise ValueError('Static pages require aligned original and bounded guest addresses')
        signatures = {}
        for target, name in self.hosts.items():
            signature = target in self.returning_hosts
            if name in signatures and signatures[name] != signature:
                raise ValueError('A host import cannot have conflicting ABIs')
            signatures[name] = signature

    def decode(self, word, pc):
        if pc in self.zero_stores:
            rn = (word >> 5) & 31
            mode, count = 2, 1
            if word & 0xffc00000 == 0x3d800000:
                offset = ((word >> 10) & 4095)*16
            elif word & 0xffe00000 == 0x3c800000:
                mode = (word >> 10)&3
                if mode == 2: raise UnsupportedInstruction('Unprivileged vector store')
                offset = signed((word >> 12)&511,9)
            else:
                mode, count = (word >> 23)&3, 2
                if mode == 0: raise UnsupportedInstruction('Non-temporal vector pair')
                offset = signed((word >> 15)&127,7)*16
            code = []
            if mode == 3: code.append(self.put(rn,self.address(rn,offset),sp=True))
            code.append(f'(local.set $base {self.address(rn,offset if mode in (0,2) else 0)})')
            code += [self.store(8,f'(i64.add (local.get $base) {const(i*8)})',const(0)) for i in range(count*2)]
            if mode == 1: code.append(self.put(rn,self.address(rn,offset),sp=True))
            return code, 'zero128-store'
        if word & 0xffffffe0 == 0x6f00e400:
            return ['(nop)'], 'zero128-immediate'
        if word & 0x9f000000 == 0x90000000:
            immediate = ((word >> 5) & 0x7ffff) << 2 | ((word >> 29) & 3)
            page = ((pc & ~4095) + signed(immediate, 21)*4096) & ((1 << 64)-1)
            if page not in self.static_pages:
                raise UnsupportedInstruction(f'Unmapped ADRP page {page:#x}')
            return [self.put(word & 31, const(self.static_pages[page]))], 'mapped-static-page'
        if word & 0x1f800000 == 0x12000000:
            width = 64 if word >> 31 else 32
            n, r, s, operation = (word >> 22)&1, (word >> 16)&63, (word >> 10)&63, (word >> 29)&3
            length = ((n << 6) | ((~s)&63)).bit_length()-1
            if operation == 3 or length < 1 or (width == 32 and n):
                raise UnsupportedInstruction('Unsupported logical immediate encoding')
            levels = (1 << length)-1
            if s & levels == levels:
                raise UnsupportedInstruction('Reserved logical immediate mask')
            element_size, rotation = 1 << length, r & levels
            mask = (1 << element_size)-1
            element = (1 << ((s & levels)+1))-1
            element = ((element >> rotation) | (element << (element_size-rotation))) & mask
            immediate = sum(element << shift for shift in range(0,width,element_size))
            expression = f'(i64.{["and","or","xor"][operation]} {self.reg((word >> 5)&31,width)} {const(immediate)})'
            return [self.put(word&31,expression,width)], 'logical-immediate'
        # ASR immediate only; other signed bitfield forms remain unsupported.
        if word & 0x7f800000 == 0x13000000:
            width = 64 if word >> 31 else 32
            r, s = (word >> 16) & 63, (word >> 10) & 63
            if ((word >> 22) & 1) != (width == 64) or s != width-1 or r >= width:
                raise UnsupportedInstruction('Unsupported signed bitfield form')
            value = self.reg((word >> 5) & 31, width)
            expression = f'(i64.shr_s {value} {const(r)})' if width == 64 else \
                f'(i64.extend_i32_u (i32.shr_s (i32.wrap_i64 {value}) (i32.const {r})))'
            return [self.put(word & 31, expression, width)], 'arithmetic-shift'
        if word & 0xffe0001f == 0xd4200000:
            return ['unreachable'], 'break-trap'
        if word & 0xfc000000 == 0x94000000:
            target = pc + signed(word & 0x3ffffff, 26)*4
            if target in self.returning_hosts:
                host = self.hosts[target]
                self.calls.append({'site':hex(pc), 'target':hex(target), 'host':host, 'abi':'context-x0-result'})
                registers = [f'x{i}' for i in range(31)]+['sp']
                # Snapshot the complete register file; the adapter may return x0
                # only. All other registers remain in their guest locals.
                return [self.put(30, const(pc+4)),
                    *[f'(i64.store offset={i*8} (local.get $ctx) (local.get ${name}))' for i,name in enumerate(registers)],
                    f'(local.set $status (call ${host} (local.get $ctx)))',
                    '(br_if $done (local.get $status))',
                    '(local.set $x0 (i64.load (local.get $ctx)))'], 'bl-host'
        return super().decode(word, pc)

    def compile(self, functions, entry):
        self.instructions, self.calls = [], []
        if entry not in functions or not functions:
            raise ValueError('Entry must name a declared guest function')
        sites, pairs = {}, set()
        self.zero_stores = {}
        for address, body in sorted(functions.items()):
            if address < 0 or address % 4 or not body or len(body) % 4 or address + len(body) > 1 << 64:
                raise ValueError('Guest functions require aligned, nonempty 64-bit instruction ranges')
            words = struct.unpack('<'+'I'*(len(body)//4),body)
            # Identify stores first; validate definitely-zero source registers
            # against the complete control-flow graph below. Other SIMD remains
            # unsupported. The existing adjacent load/copy idiom stays separate.
            for i, word in enumerate(words):
                if word & 0xffe00000 == 0x3c800000 or word & 0xffc00000 == 0x3d800000:
                    is_copy = (word & 0xffc00000 == 0x3d800000 and i and
                               words[i-1] & 0xffe00c1f == 0x3cc00000 | (word&31))
                    if not is_copy: self.zero_stores[address+i*4] = [word&31]
                elif word & 0x3a000000 == 0x28000000 and word & (1<<26) and word >> 30 == 2 and not word & (1<<22):
                    self.zero_stores[address+i*4] = [word&31,(word>>10)&31]
            for i, word in enumerate(words):
                pc = address+i*4
                if pc in sites:
                    raise ValueError('Overlapping guest functions')
                sites[pc] = word
                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 and pc not in self.zero_stores:
                    if not i or words[i-1] & 0xffe00c1f != 0x3cc00000 | (word & 31):
                        raise UnsupportedInstruction('128-bit store requires adjacent matching load')
                    pairs.add(pc)
        if set(sites) & set(self.hosts):
            raise ValueError('Guest and host entry points overlap')
        addresses = sorted(sites)
        indices = {pc:i for i,pc in enumerate(addresses)}
        guest_calls = {}
        for pc, word in sites.items():
            if word & 0xfc000000 == 0x94000000:
                target = pc + ((word & 0x3ffffff) - (0x4000000 if word & 0x2000000 else 0))*4
                if target not in self.hosts:
                    if target not in functions:
                        raise UnsupportedInstruction(f'undeclared guest BL target {target:#x}')
                    if pc+4 not in sites:
                        raise UnsupportedInstruction('Guest call has no declared continuation')
                    guest_calls[pc] = target
        returns = sorted(pc+4 for pc in guest_calls)
        successors, kinds, nodes = {}, {}, []
        for pc in addresses:
            word = sites[pc]
            branch = self.branch(word, pc)
            if branch:
                target, condition = branch
                edges = [target] + ([pc+4] if condition is not None else [])
                next_pc = f'(i32.const {indices.get(target, -1)})'
                if condition is not None:
                    next_pc = f'(select {next_pc} (i32.const {indices.get(pc+4, -1)}) {condition})'
                emitted = [f'(local.set $pc {next_pc})','(br $dispatch)']
                kind = 'branch-immediate'
                if target in pairs:
                    raise UnsupportedInstruction('Control flow enters a 128-bit copy halfway')
            elif pc in guest_calls:
                target = guest_calls[pc]
                emitted = [self.put(30,const(pc+4)),
                    f'(local.set $pc (i32.const {indices[target]}))','(br $dispatch)']
                edges, kind = [target], 'bl-guest'
                self.calls.append({'site':hex(pc),'target':hex(target),'kind':'guest'})
            elif word == 0xd65f03c0:
                emitted = ['(br_if $done (i64.eq (local.get $x30) (local.get $exit_lr)))']
                emitted += [f'(if (i64.eq (local.get $x30) {const(target)}) (then '
                    f'(local.set $pc (i32.const {indices[target]})) (br $dispatch)))' for target in returns]
                emitted += ['(local.set $status (i32.const 4))','(br $done)']
                edges, kind = returns, 'ret'
            else:
                try:
                    emitted, kind = self.decode(word,pc)
                except UnsupportedInstruction as error:
                    raise UnsupportedInstruction(f'{pc:#x}: {error}') from error
                emitted += [f'(local.set $pc (i32.const {indices.get(pc+4,-1)}))','(br $dispatch)']
                edges = [] if kind == 'break-trap' else [pc+4]
            for target in edges:
                if target not in sites:
                    raise UnsupportedInstruction(f'Control flow outside declared code: {target:#x}')
            successors[pc], kinds[pc] = edges, kind
            self.instructions.append({'address':hex(pc),'word':f'{word:08x}','kind':kind})
            nodes.append(f'(if (i32.eq (local.get $pc) (i32.const {indices[pc]})) (then\n'+ '\n'.join(emitted)+'\n))')
        reachable, pending = set(), [entry]
        while pending:
            pc = pending.pop()
            if pc not in reachable:
                reachable.add(pc)
                pending.extend(successors[pc])
        if reachable != set(sites):
            raise UnsupportedInstruction('Unreachable declared code requires explicit classification')
        if not any(kind=='ret' for kind in kinds.values()):
            raise UnsupportedInstruction('Program has no RET')
        # Conservative flag analysis includes every possible declared return edge.
        predecessors = {pc:[] for pc in sites}
        for pc, edges in successors.items():
            for target in edges: predecessors[target].append(pc)
        incoming = {pc: pc!=entry for pc in sites}
        outgoing = dict(incoming)
        changed = True
        while changed:
            changed = False
            for pc in addresses:
                value = pc!=entry and all(outgoing[p] for p in predecessors[pc])
                after = True if kinds[pc]=='arithmetic-flags' else False if kinds[pc]=='bl-host' else value
                if (incoming[pc],outgoing[pc])!=(value,after):
                    incoming[pc],outgoing[pc],changed = value,after,True
        for pc,word in sites.items():
            condition = word & 15 if word & 0xff000010 == 0x54000000 else \
                (word>>12)&15 if kinds[pc]=='select-flags' else 14
            if condition<14 and not incoming[pc]:
                raise UnsupportedInstruction(f'Conditional operation reads undefined NZCV at {pc:#x}')
        # Definitely-zero facts are intersected at every merge. Calls clear all
        # facts (the ABI preserves no complete 128-bit vector register). An
        # adjacent copy load also overwrites its destination vector register.
        zero_in = {pc:0 if pc==entry else 0xffffffff for pc in sites}
        zero_out = dict(zero_in)
        changed = True
        while changed:
            changed = False
            for pc in addresses:
                before = 0 if pc==entry else 0xffffffff
                for parent in predecessors[pc]: before &= zero_out[parent]
                after = before
                if kinds[pc] in ('bl-host','bl-guest'): after = 0
                elif kinds[pc]=='zero128-immediate': after |= 1 << (sites[pc]&31)
                elif kinds[pc]=='copy128-load': after &= ~(1 << (sites[pc]&31))
                if (zero_in[pc],zero_out[pc])!=(before,after):
                    zero_in[pc],zero_out[pc],changed = before,after,True
        for pc, registers in self.zero_stores.items():
            if any(not zero_in[pc] & (1 << register) for register in registers):
                raise UnsupportedInstruction(f'Vector store reads undefined/nonzero Q register at {pc:#x}')
        registers = [f'x{i}' for i in range(31)]+['sp']
        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))
        returning_names = {self.hosts[target] for target in self.returning_hosts}
        imports = '\n'.join(f'(import "env" "{name}" (func ${name} (param {"i32" if name in returning_names else "i64 i64 i64"}) (result i32)))'
                            for name in sorted(set(self.hosts.values())))
        # The external return sentinel must not alias any declared instruction.
        sentinel_checks = '\n'.join(f'(if (i64.eq (local.get $exit_lr) {const(pc)}) (then (return (i32.const 2))))'
                                    for pc in addresses)
        return f'''(module
{imports}
(memory (export "memory") {self.memory_pages} {self.memory_pages})
(func $addr (param $value i64) (param $size i64) (result i32)
 (if (i64.gt_u (local.get $value) (i64.sub (i64.const {self.memory_pages*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 $exit_lr i64)
(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)
(if (i32.gt_u (local.get $ctx) (i32.const {self.memory_pages*65536-256})) (then (return (i32.const 2))))
{loads}
(local.set $exit_lr (local.get $x30))
{sentinel_checks}
(local.set $pc (i32.const {indices[entry]}))
(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)))
'''