from __future__ import annotations

from collections.abc import Iterable

from .formatting import (
    bit_comment,
    control_reg,
    disp_text,
    h8,
    h16,
    h24,
    label_or_h,
    reg_list,
    s8,
    s16,
    write_comment,
)
from .model import EA, Instruction
from .rom import DecodeError, Rom
from .tables import BRANCH_NAMES, IO_REGISTERS


class H8536Decoder:
    def __init__(self, rom: Rom, br: int | None = None, labels: dict[int, str] | None = None) -> None:
        self.rom = rom
        self.br = br
        self.labels = labels if labels is not None else {}

    def symbol(self, address: int) -> str:
        if address in self.labels:
            return self.labels[address]
        if address in IO_REGISTERS:
            return IO_REGISTERS[address]
        return h16(address)

    def mem(self, address: int) -> str:
        return "@" + self.symbol(address)

    def short_abs_address(self, low: int) -> int | None:
        if self.br is None:
            return None
        return ((self.br & 0xFF) << 8) | (low & 0xFF)

    def short_abs(self, low: int) -> str:
        address = self.short_abs_address(low)
        if address is None:
            return f"@BR:{h8(low)}"
        return self.mem(address)

    def decode_ea(self, address: int) -> EA:
        b = self.rom.u8(address)
        if 0xA0 <= b <= 0xAF:
            size = "W" if b & 0x08 else "B"
            reg = b & 0x07
            return EA(f"R{reg}", "reg", size, 1, reg=reg)
        if 0xB0 <= b <= 0xBF:
            size = "W" if b & 0x08 else "B"
            reg = b & 0x07
            return EA(f"@-R{reg}", "predec", size, 1, reg=reg)
        if 0xC0 <= b <= 0xCF:
            size = "W" if b & 0x08 else "B"
            reg = b & 0x07
            return EA(f"@R{reg}+", "postinc", size, 1, reg=reg)
        if 0xD0 <= b <= 0xDF:
            size = "W" if b & 0x08 else "B"
            reg = b & 0x07
            return EA(f"@R{reg}", "indirect", size, 1, reg=reg)
        if 0xE0 <= b <= 0xEF:
            size = "W" if b & 0x08 else "B"
            reg = b & 0x07
            disp = self.rom.u8(address + 1)
            return EA(f"@({disp_text(s8(disp))},R{reg})", "disp8", size, 2, reg=reg, value=s8(disp))
        if 0xF0 <= b <= 0xFF:
            size = "W" if b & 0x08 else "B"
            reg = b & 0x07
            disp = self.rom.u16(address + 1)
            return EA(f"@({disp_text(s16(disp))},R{reg})", "disp16", size, 3, reg=reg, value=s16(disp))
        if b == 0x04:
            value = self.rom.u8(address + 1)
            return EA(f"#{h8(value)}", "imm", "B", 2, value=value)
        if b == 0x0C:
            value = self.rom.u16(address + 1)
            return EA(f"#{h16(value)}", "imm", "W", 3, value=value)
        if b in (0x05, 0x0D):
            size = "W" if b == 0x0D else "B"
            low = self.rom.u8(address + 1)
            full = self.short_abs_address(low)
            text = self.short_abs(low)
            return EA(text, "abs8", size, 2, value=low, address=full)
        if b in (0x15, 0x1D):
            size = "W" if b == 0x1D else "B"
            full = self.rom.u16(address + 1)
            return EA(self.mem(full), "abs16", size, 3, address=full)
        raise DecodeError(f"not an EA byte: {h8(b)}")

    def decode(self, address: int) -> Instruction:
        try:
            b = self.rom.u8(address)
        except DecodeError:
            return Instruction(address, b"", ".eof", kind="invalid", fallthrough=False, valid=False)

        if is_ea_byte(b):
            try:
                return self.decode_general(address)
            except DecodeError:
                pass

        try:
            return self.decode_direct(address)
        except DecodeError:
            raw = self.rom.slice(address, 1)
            return Instruction(address, raw, ".db", h8(raw[0]), kind="invalid", fallthrough=False, valid=False)

    def decode_general(self, address: int) -> Instruction:
        ea = self.decode_ea(address)
        op_addr = address + ea.length
        op = self.rom.u8(op_addr)
        end = op_addr + 1
        mnemonic = ""
        operands = ""
        comment = ""
        writes_br = False
        br_value: int | None = None

        source_to_reg = {
            0x20: "ADD:G",
            0x28: "ADDS",
            0x30: "SUB",
            0x38: "SUBS",
            0x40: "OR",
            0x50: "AND",
            0x60: "XOR",
            0x70: "CMP:G",
            0x80: "MOV:G",
            0xA0: "ADDX",
            0xA8: "MULXU",
            0xB0: "SUBX",
            0xB8: "DIVXU",
        }

        base = op & 0xF8
        rd = op & 0x07

        if base in source_to_reg:
            mnemonic = f"{source_to_reg[base]}.{ea.size}"
            operands = f"{ea.text}, R{rd}"
        elif base == 0x90:
            rs = op & 0x07
            if ea.mode == "reg":
                mnemonic = "XCH.W"
                operands = f"R{ea.reg}, R{rs}"
            else:
                mnemonic = f"MOV:G.{ea.size}"
                operands = f"R{rs}, {ea.text}"
                comment = write_comment(ea, None)
        elif op in (0x10, 0x11, 0x12) and ea.mode == "reg":
            names = {0x10: "SWAP.B", 0x11: "EXTS.B", 0x12: "EXTU.B"}
            mnemonic = names[op]
            operands = ea.text
        elif op in (0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F):
            unary = {
                0x13: "CLR",
                0x14: "NEG",
                0x15: "NOT",
                0x16: "TST",
                0x17: "TAS",
                0x18: "SHAL",
                0x19: "SHAR",
                0x1A: "SHLL",
                0x1B: "SHLR",
                0x1C: "ROTL",
                0x1D: "ROTR",
                0x1E: "ROTXL",
                0x1F: "ROTXR",
            }
            size = "B" if op == 0x17 else ea.size
            mnemonic = f"{unary[op]}.{size}"
            operands = ea.text
        elif op in (0x08, 0x09, 0x0C, 0x0D):
            q = {0x08: 1, 0x09: 2, 0x0C: -1, 0x0D: -2}[op]
            mnemonic = f"ADD:Q.{ea.size}"
            operands = f"#{q}, {ea.text}"
            comment = write_comment(ea, None)
        elif op in (0x06, 0x07):
            if op == 0x06:
                imm = self.rom.u8(end)
                end += 1
                imm_text = h8(imm)
                size = ea.size
            else:
                imm = self.rom.u16(end)
                end += 2
                imm_text = h16(imm)
                size = "W"
            mnemonic = f"MOV:G.{size}"
            operands = f"#{imm_text}, {ea.text}"
            comment = write_comment(ea, imm)
        elif op in (0x04, 0x05):
            if op == 0x04:
                imm = self.rom.u8(end)
                end += 1
                imm_text = h8(imm)
                size = "B"
            else:
                imm = self.rom.u16(end)
                end += 2
                imm_text = h16(imm)
                size = "W"
            mnemonic = f"CMP:G.{size}"
            operands = f"#{imm_text}, {ea.text}"
        elif op == 0x00:
            ext = self.rom.u8(end)
            end += 1
            ext_base = ext & 0xF8
            if ext_base == 0x80:
                mnemonic = "MOVFPE.B"
                operands = f"{ea.text}, R{ext & 7}"
            elif ext_base == 0x90:
                mnemonic = "MOVTPE.B"
                operands = f"R{ext & 7}, {ea.text}"
                comment = write_comment(ea, None)
            else:
                raise DecodeError(f"unknown extended general op {h8(ext)} after {h8(op)}")
        elif (op & 0xF0) in (0xC0, 0xD0, 0xE0, 0xF0):
            bit_ops = {0xC0: "BSET", 0xD0: "BCLR", 0xE0: "BNOT", 0xF0: "BTST"}
            mnemonic = f"{bit_ops[op & 0xF0]}.{ea.size}"
            operands = f"#{op & 0x0F}, {ea.text}"
            if op & 0xF0 in (0xC0, 0xD0, 0xE0):
                comment = bit_comment(mnemonic, ea, op & 0x0F)
        elif base in (0x48, 0x58, 0x68, 0x78) and ea.mode != "imm":
            bit_ops = {0x48: "BSET", 0x58: "BCLR", 0x68: "BNOT", 0x78: "BTST"}
            mnemonic = f"{bit_ops[base]}.{ea.size}"
            operands = f"R{rd}, {ea.text}"
        elif base == 0x88:
            mnemonic = f"LDC.{ea.size}"
            operands = f"{ea.text}, {control_reg(rd, ea.size)}"
            if rd == 1 and ea.size == "B":
                writes_br = True
                br_value = ea.value if ea.mode == "imm" else None
                if br_value is not None:
                    comment = f"BR = {h8(br_value)}; @aa:8 upper address byte"
        elif base == 0x98:
            mnemonic = f"STC.{ea.size}"
            operands = f"{control_reg(rd, ea.size)}, {ea.text}"
            comment = write_comment(ea, None)
        elif ea.mode == "imm" and base in (0x48, 0x58, 0x68):
            ctrl_ops = {0x48: "ORC", 0x58: "ANDC", 0x68: "XORC"}
            mnemonic = f"{ctrl_ops[base]}.{ea.size}"
            operands = f"{ea.text}, {control_reg(rd, ea.size)}"
        else:
            raise DecodeError(f"unknown general op {h8(op)} after {h8(self.rom.u8(address))}")

        raw = self.rom.slice(address, end - address)
        references = [ea.address] if ea.address is not None else []
        return Instruction(
            address,
            raw,
            mnemonic,
            operands,
            comment=comment,
            references=references,
            writes_br=writes_br,
            br_value=br_value,
        )

    def decode_direct(self, address: int) -> Instruction:
        b = self.rom.u8(address)

        if b == 0x00:
            return self.ins(address, 1, "NOP")
        if b == 0x02:
            mask = self.rom.u8(address + 1)
            return self.ins(address, 2, "LDM.W", f"@SP+, {reg_list(mask)}")
        if b == 0x12:
            mask = self.rom.u8(address + 1)
            return self.ins(address, 2, "STM.W", f"{reg_list(mask)}, @-SP")
        if b == 0x03:
            page = self.rom.u8(address + 1)
            pc = self.rom.u16(address + 2)
            return self.ins(address, 4, "PJSR", f"@{h24((page << 16) | pc)}", "call", [pc])
        if b == 0x13:
            page = self.rom.u8(address + 1)
            pc = self.rom.u16(address + 2)
            return self.ins(address, 4, "PJMP", f"@{h24((page << 16) | pc)}", "jump", [pc], False)
        if b in (0x01, 0x06, 0x07):
            op2 = self.rom.u8(address + 1)
            if 0xB8 <= op2 <= 0xBF:
                disp = s8(self.rom.u8(address + 2))
                target = (address + 3 + disp) & 0xFFFF
                cond = {0x01: "F", 0x06: "NE", 0x07: "EQ"}[b]
                return self.ins(
                    address,
                    3,
                    f"SCB/{cond}",
                    f"R{op2 & 7}, {label_or_h(target, self.labels)}",
                    "branch",
                    [target],
                    True,
                )
        if b == 0x08:
            op2 = self.rom.u8(address + 1)
            if (op2 & 0xF0) == 0x10:
                return self.ins(address, 2, "TRAPA", f"#{op2 & 0x0F}")
        if b == 0x09:
            return self.ins(address, 1, "TRAP/VS")
        if b == 0x0A:
            return self.ins(address, 1, "RTE", kind="rte", fallthrough=False)
        if b == 0x0E:
            disp = s8(self.rom.u8(address + 1))
            target = (address + 2 + disp) & 0xFFFF
            return self.ins(address, 2, "BSR", label_or_h(target, self.labels), "call", [target])
        if b == 0x0F:
            return self.ins(address, 1, "UNLK", "FP")
        if b == 0x10:
            target = self.rom.u16(address + 1)
            return self.ins(address, 3, "JMP", f"@{label_or_h(target, self.labels)}", "jump", [target], False)
        if b == 0x11:
            return self.decode_11(address)
        if b == 0x14:
            value = self.rom.u8(address + 1)
            return self.ins(address, 2, "RTD", f"#{h8(value)}", "return", fallthrough=False)
        if b == 0x17:
            value = self.rom.u8(address + 1)
            return self.ins(address, 2, "LINK", f"FP, #{h8(value)}")
        if b == 0x18:
            target = self.rom.u16(address + 1)
            return self.ins(address, 3, "JSR", f"@{label_or_h(target, self.labels)}", "call", [target])
        if b == 0x19:
            return self.ins(address, 1, "RTS", kind="return", fallthrough=False)
        if b == 0x1A:
            return self.ins(address, 1, "SLEEP", kind="sleep", fallthrough=False)
        if b == 0x1C:
            value = self.rom.u16(address + 1)
            return self.ins(address, 3, "RTD", f"#{h16(value)}", "return", fallthrough=False)
        if b == 0x1E:
            disp = s16(self.rom.u16(address + 1))
            target = (address + 3 + disp) & 0xFFFF
            return self.ins(address, 3, "BSR", label_or_h(target, self.labels), "call", [target])
        if b == 0x1F:
            value = self.rom.u16(address + 1)
            return self.ins(address, 3, "LINK", f"FP, #{h16(value)}")
        if 0x20 <= b <= 0x2F:
            cond = b & 0x0F
            disp = s8(self.rom.u8(address + 1))
            target = (address + 2 + disp) & 0xFFFF
            fallthrough = BRANCH_NAMES[cond] not in ("BRA",)
            kind = "jump" if BRANCH_NAMES[cond] == "BRA" else "branch"
            return self.ins(address, 2, BRANCH_NAMES[cond], label_or_h(target, self.labels), kind, [target], fallthrough)
        if 0x30 <= b <= 0x3F:
            cond = b & 0x0F
            disp = s16(self.rom.u16(address + 1))
            target = (address + 3 + disp) & 0xFFFF
            fallthrough = BRANCH_NAMES[cond] not in ("BRA",)
            kind = "jump" if BRANCH_NAMES[cond] == "BRA" else "branch"
            return self.ins(address, 3, BRANCH_NAMES[cond], label_or_h(target, self.labels), kind, [target], fallthrough)
        if 0x40 <= b <= 0x47:
            rd = b & 7
            imm = self.rom.u8(address + 1)
            return self.ins(address, 2, "CMP:E", f"#{h8(imm)}, R{rd}")
        if 0x48 <= b <= 0x4F:
            rd = b & 7
            imm = self.rom.u16(address + 1)
            return self.ins(address, 3, "CMP:I", f"#{h16(imm)}, R{rd}")
        if 0x50 <= b <= 0x57:
            rd = b & 7
            imm = self.rom.u8(address + 1)
            return self.ins(address, 2, "MOV:E.B", f"#{h8(imm)}, R{rd}")
        if 0x58 <= b <= 0x5F:
            rd = b & 7
            imm = self.rom.u16(address + 1)
            return self.ins(address, 3, "MOV:I.W", f"#{h16(imm)}, R{rd}")
        if 0x60 <= b <= 0x67:
            rd = b & 7
            low = self.rom.u8(address + 1)
            ref = self.short_abs_address(low)
            return self.ins(
                address,
                2,
                "MOV:L.B",
                f"{self.short_abs(low)}, R{rd}",
                references=[ref] if ref is not None else None,
            )
        if 0x68 <= b <= 0x6F:
            rd = b & 7
            low = self.rom.u8(address + 1)
            ref = self.short_abs_address(low)
            return self.ins(
                address,
                2,
                "MOV:L.W",
                f"{self.short_abs(low)}, R{rd}",
                references=[ref] if ref is not None else None,
            )
        if 0x70 <= b <= 0x77:
            rs = b & 7
            low = self.rom.u8(address + 1)
            ref = self.short_abs_address(low)
            ea = EA(self.short_abs(low), "abs8", "B", 2, value=low, address=ref)
            return self.ins(
                address,
                2,
                "MOV:S.B",
                f"R{rs}, {ea.text}",
                references=[ref] if ref is not None else None,
                comment=write_comment(ea, None),
            )
        if 0x78 <= b <= 0x7F:
            rs = b & 7
            low = self.rom.u8(address + 1)
            ref = self.short_abs_address(low)
            ea = EA(self.short_abs(low), "abs8", "W", 2, value=low, address=ref)
            return self.ins(
                address,
                2,
                "MOV:S.W",
                f"R{rs}, {ea.text}",
                references=[ref] if ref is not None else None,
                comment=write_comment(ea, None),
            )
        if 0x80 <= b <= 0x87:
            rd = b & 7
            disp = s8(self.rom.u8(address + 1))
            return self.ins(address, 2, "MOV:F.B", f"@({disp_text(disp)},R6), R{rd}")
        if 0x88 <= b <= 0x8F:
            rd = b & 7
            disp = s8(self.rom.u8(address + 1))
            return self.ins(address, 2, "MOV:F.W", f"@({disp_text(disp)},R6), R{rd}")
        if 0x90 <= b <= 0x97:
            rs = b & 7
            disp = s8(self.rom.u8(address + 1))
            return self.ins(address, 2, "MOV:F.B", f"R{rs}, @({disp_text(disp)},R6)")
        if 0x98 <= b <= 0x9F:
            rs = b & 7
            disp = s8(self.rom.u8(address + 1))
            return self.ins(address, 2, "MOV:F.W", f"R{rs}, @({disp_text(disp)},R6)")
        raise DecodeError(f"unknown direct opcode {h8(b)}")

    def decode_11(self, address: int) -> Instruction:
        op = self.rom.u8(address + 1)
        if op == 0x14:
            value = self.rom.u8(address + 2)
            return self.ins(address, 3, "PRTD", f"#{h8(value)}", "return", fallthrough=False)
        if op == 0x19:
            return self.ins(address, 2, "PRTS", kind="return", fallthrough=False)
        if op == 0x1C:
            value = self.rom.u16(address + 2)
            return self.ins(address, 4, "PRTD", f"#{h16(value)}", "return", fallthrough=False)
        if 0xC0 <= op <= 0xC7:
            return self.ins(address, 2, "PJMP", f"@R{op & 7}", "jump", fallthrough=False)
        if 0xC8 <= op <= 0xCF:
            return self.ins(address, 2, "PJSR", f"@R{op & 7}", "call")
        if 0xD0 <= op <= 0xD7:
            return self.ins(address, 2, "JMP", f"@R{op & 7}", "jump", fallthrough=False)
        if 0xD8 <= op <= 0xDF:
            return self.ins(address, 2, "JSR", f"@R{op & 7}", "call")
        if 0xE0 <= op <= 0xE7:
            disp = s8(self.rom.u8(address + 2))
            return self.ins(address, 3, "JMP", f"@({disp_text(disp)},R{op & 7})", "jump", fallthrough=False)
        if 0xE8 <= op <= 0xEF:
            disp = s8(self.rom.u8(address + 2))
            return self.ins(address, 3, "JSR", f"@({disp_text(disp)},R{op & 7})", "call")
        if 0xF0 <= op <= 0xF7:
            disp = s16(self.rom.u16(address + 2))
            return self.ins(address, 4, "JMP", f"@({disp_text(disp)},R{op & 7})", "jump", fallthrough=False)
        if 0xF8 <= op <= 0xFF:
            disp = s16(self.rom.u16(address + 2))
            return self.ins(address, 4, "JSR", f"@({disp_text(disp)},R{op & 7})", "call")
        raise DecodeError(f"unknown 0x11 opcode {h8(op)}")

    def ins(
        self,
        address: int,
        size: int,
        mnemonic: str,
        operands: str = "",
        kind: str = "normal",
        targets: Iterable[int] | None = None,
        fallthrough: bool = True,
        references: Iterable[int] | None = None,
        comment: str = "",
        writes_br: bool = False,
        br_value: int | None = None,
    ) -> Instruction:
        return Instruction(
            address,
            self.rom.slice(address, size),
            mnemonic,
            operands,
            kind=kind,
            targets=list(targets or []),
            fallthrough=fallthrough,
            comment=comment,
            references=list(references or []),
            writes_br=writes_br,
            br_value=br_value,
        )


def is_ea_byte(value: int) -> bool:
    return value in (0x04, 0x05, 0x0C, 0x0D, 0x15, 0x1D) or 0xA0 <= value <= 0xFF
