"""A tiny Z80 machine-code emitter (just what the ColecoVision/Adam viewer needs).

No Z80 assembler is installed, so -- as with the TMS9900 and 6809 emitters -- we
emit opcodes directly.  Little-endian 16-bit operands; 8-bit relative jumps are
backpatched via labels.
"""

from __future__ import annotations


class Asm:
    def __init__(self, base: int):
        self.base = base
        self.code = bytearray()
        self.labels: dict[str, int] = {}
        self._fix: list[tuple[int, str]] = []      # (pos of rel byte, label)

    def pos(self) -> int:
        return self.base + len(self.code)

    def label(self, name: str):
        self.labels[name] = self.pos()

    def _b(self, *bs):
        self.code += bytes(b & 0xFF for b in bs)

    def _w(self, v):                                # 16-bit, little-endian
        self.code += bytes([v & 0xFF, (v >> 8) & 0xFF])

    # ---- 8-bit loads (immediate) ----
    def ld_a(self, n):  self._b(0x3E, n)
    def ld_b(self, n):  self._b(0x06, n)
    def ld_c(self, n):  self._b(0x0E, n)
    def ld_d(self, n):  self._b(0x16, n)
    def ld_e(self, n):  self._b(0x1E, n)

    # ---- 16-bit loads (immediate) ----
    def ld_hl(self, nn): self._b(0x21); self._w(nn)
    def ld_de(self, nn): self._b(0x11); self._w(nn)
    def ld_bc(self, nn): self._b(0x01); self._w(nn)

    # ---- memory <-> A (extended addressing) ----
    def ld_a_mem(self, addr):  self._b(0x3A); self._w(addr)     # LD A,(nn)
    def ld_mem_a(self, addr):  self._b(0x32); self._w(addr)     # LD (nn),A

    # ---- register moves ----
    def ld_a_b(self): self._b(0x78)
    def ld_a_c(self): self._b(0x79)
    def ld_a_d(self): self._b(0x7A)
    def ld_a_e(self): self._b(0x7B)
    def ld_a_h(self): self._b(0x7C)
    def ld_a_l(self): self._b(0x7D)
    def ld_b_a(self): self._b(0x47)
    def ld_c_a(self): self._b(0x4F)
    def ld_a_hl(self): self._b(0x7E)        # LD A,(HL)
    def ld_hl_a(self): self._b(0x77)        # LD (HL),A

    # ---- I/O ----
    def out_n_a(self, n): self._b(0xD3, n)  # OUT (n),A
    def in_a_n(self, n):  self._b(0xDB, n)  # IN A,(n)

    # ---- arithmetic / logic ----
    def inc_a(self):  self._b(0x3C)
    def dec_a(self):  self._b(0x3D)
    def dec_c(self):  self._b(0x0D)
    def inc_hl(self): self._b(0x23)
    def dec_hl(self): self._b(0x2B)
    def inc_de(self): self._b(0x13)
    def dec_de(self): self._b(0x1B)
    def dec_bc(self): self._b(0x0B)
    def or_a(self):   self._b(0xB7)         # OR A (set flags from A)
    def or_c(self):   self._b(0xB1)
    def or_e(self):   self._b(0xB3)
    def and_n(self, n): self._b(0xE6, n)
    def cp_n(self, n):  self._b(0xFE, n)
    def add_hl_de(self): self._b(0x19)

    # ---- block ops ----
    def ldir(self): self._b(0xED, 0xB0)     # (HL)->(DE), BC times

    # ---- control ----
    def di(self):   self._b(0xF3)
    def ei(self):   self._b(0xFB)
    def ret(self):  self._b(0xC9)
    def retn(self): self._b(0xED, 0x45)
    def nop(self):  self._b(0x00)

    def jp(self, addr):     self._b(0xC3); self._w(addr)
    def jp_label(self, lbl):
        self._b(0xC3); self._fix.append((len(self.code), lbl)); self._w(0)

    def _jr(self, op, lbl):
        self._b(op); self._fix.append((len(self.code), lbl)); self._b(0)

    def jr(self, lbl):    self._jr(0x18, lbl)
    def jr_nz(self, lbl): self._jr(0x20, lbl)
    def jr_z(self, lbl):  self._jr(0x28, lbl)
    def djnz(self, lbl):  self._jr(0x10, lbl)

    def resolve(self) -> bytes:
        for pos, lbl in self._fix:
            target = self.labels[lbl]
            opcode = self.code[pos - 1]                  # the byte just before the operand
            if opcode == 0xC3:                           # JP nn (absolute, 2 bytes)
                self.code[pos] = target & 0xFF
                self.code[pos + 1] = (target >> 8) & 0xFF
            else:                                        # JR / DJNZ (relative, 1 byte)
                rel = target - (self.base + pos + 1)
                if not -128 <= rel <= 127:
                    raise ValueError(f"relative jump to {lbl} out of range ({rel})")
                self.code[pos] = rel & 0xFF
        return bytes(self.code)