from dataclasses import dataclass
import struct
from typing import Any, Callable, Iterable, Literal, Optional


class Label:
    "A label in the program."

    name: Optional[str]

    def __init__(self, name=None):
        self.name = name


@dataclass
class Insn:
    "An instruction that has not yet been relocated."

    len: int
    fmt: str
    relocate: Callable[[int, dict[Label, int]], Iterable[Any]]


@dataclass
class DerefReg:
    "An indirect reference through a register."

    reg: "Reg"


@dataclass
class Reg:
    "A register."

    num: int
    derefable: bool = False

    def __str__(self):
        return f"R{self.num}"

    def deref(self) -> DerefReg:
        if not self.derefable:
            raise Exception("Cannot deref through {self}")
        return DerefReg(self)


@dataclass
class SFR:
    "A reference to one of the 256 bytes of iRAM (including the SFRs, hence this type's name)."

    num: int


class Asm:
    """
    The interface to the assembler. This class has a method for each of the
    8051's instructions.
    """

    R0 = Reg(0, derefable=True)
    R1 = Reg(1, derefable=True)
    R2 = Reg(2)
    R3 = Reg(3)
    R4 = Reg(4)
    R5 = Reg(5)
    R6 = Reg(6)
    R7 = Reg(7)

    def __init__(self):
        insns: list["Insn"] = []
        self.addr = 0x0000
        self.insns = insns
        self.labels = {}

    def assemble(self) -> bytes:
        out = b""
        addr = 0
        for insn in self.insns:
            out += struct.pack(insn.fmt, *insn.relocate(addr, self.labels))
            addr += insn.len
        return out

    def deflabel(self, label: Label) -> None:
        if label in self.labels:
            raise Exception(
                f"[{hex(self.addr)}] Attempting to redefine the label {label}"
            )
        self.labels[label] = self.addr

    def ACALL(self, label: Label) -> Insn:
        """
        Absolute Call. Pushes the address of the next instruction to the stack,
        then sets the low 11 bits of the program counter to match the label.

        The assembler checks that the top 5 bits match as well.
        """

        def relocate(addr: int, labels: dict[Label, int]) -> Iterable[int]:
            label_addr = labels[label]
            if (label_addr >> 11) != ((addr + 2) >> 1):
                raise Exception(
                    f"[{hex(self.addr)}] ACALL to out-of-range label {label}"
                )
            return (0x11 | ((label_addr >> 8) & 7), label_addr & 0xFF)

        insn = Insn(2, "cc", relocate)
        self.insns.append(insn)
        return insn

    def ADD(self, arg: int | SFR | DerefReg | Reg) -> Insn:
        """
        Add Accumulator. Adds the argument to the accumulator register (SFR 0xE0).
        """

        match arg:
            case int(lit):
                return Insn(2, "cc", lambda _addr, _labels: (0x24, lit))
            case SFR(addr):
                return Insn(2, "cc", lambda _addr, _labels: (0x25, addr))
            case DerefReg(Reg(reg)):
                return Insn(2, "c", lambda _addr, _labels: (0x26 + reg,))
            case Reg(reg):
                return Insn(2, "c", lambda _addr, _labels: (0x28 + reg,))

    def ADDC(self, arg: int | SFR | DerefReg | Reg) -> Insn:
        """
        Add Accumulator With Carry. Adds the argument plus the carry bit to the accumulator register (SFR 0xE0).
        """

        match arg:
            case int(lit):
                return Insn(2, "cc", lambda _addr, _labels: (0x34, lit))
            case SFR(addr):
                return Insn(2, "cc", lambda _addr, _labels: (0x35, addr))
            case DerefReg(Reg(reg)):
                return Insn(2, "c", lambda _addr, _labels: (0x36 + reg,))
            case Reg(reg):
                return Insn(2, "c", lambda _addr, _labels: (0x38 + reg,))

    def AJMP(self, label: Label) -> Insn:
        """
        Absolute Jump. Sets the low 11 bits of the program counter to match the
        label.

        The assembler checks that the top 5 bits match as well.
        """

        def relocate(addr: int, labels: dict[Label, int]) -> Iterable[int]:
            label_addr = labels[label]
            if (label_addr >> 11) != ((addr + 2) >> 1):
                raise Exception(
                    f"[{hex(self.addr)}] AJMP to out-of-range label {label}"
                )
            return (0x01 | ((label_addr >> 8) & 7), label_addr & 0xFF)

        insn = Insn(2, "cc", relocate)
        self.insns.append(insn)
        return insn

    def ANL(self, dst: SFR | Literal["A", "C"], src: Literal["A"] | int): ...

    # ANL - Bitwise AND
    # CJNE - Compare and Jump if Not Equal
    # CLR - Clear Register
    # CPL - Complement Register
    # DA - Decimal Adjust
    # DEC - Decrement Register
    # DIV - Divide Accumulator by B
    # DJNZ - Decrement Register and Jump if Not Zero
    # INC - Increment Register
    # JB - Jump if Bit Set
    # JBC - Jump if Bit Set and Clear Bit
    # JC - Jump if Carry Set
    # JMP - Jump to Address
    # JNB - Jump if Bit Not Set
    # JNC - Jump if Carry Not Set
    # JNZ - Jump if Accumulator Not Zero
    # JZ - Jump if Accumulator Zero
    # LCALL - Long Call
    # LJMP - Long Jump
    # MOV - Move Memory
    # MOVC - Move Code Memory
    # MOVX - Move Extended Memory
    # MUL - Multiply Accumulator by B
    # NOP - No Operation
    # ORL - Bitwise OR
    # POP - Pop Value From Stack
    # PUSH - Push Value Onto Stack
    # RET - Return From Subroutine
    # RETI - Return From Interrupt
    # RL - Rotate Accumulator Left
    # RLC - Rotate Accumulator Left Through Carry
    # RR - Rotate Accumulator Right
    # RRC - Rotate Accumulator Right Through Carry
    # SETB - Set Bit
    # SJMP - Short Jump
    # SUBB - Subtract From Accumulator With Borrow
    # SWAP - Swap Accumulator Nibbles
    # XCH - Exchange Bytes
    # XCHD - Exchange Digits
    # XRL - Bitwise Exclusive OR