"""Opcode mnemonic mapping and arity classification for OR1 assembly.

This module provides:
- MNEMONIC_TO_OP: Maps assembly mnemonic strings to ALUOp or MemOp enum values
- OP_TO_MNEMONIC: Reverse mapping for serialization (handles IntEnum value collisions)
- MONADIC_OPS: Set of opcodes that are always monadic
- is_monadic() and is_dyadic(): Functions to check operand arity
"""

from typing import Optional, Union
from cm_inst import ArithOp, LogicOp, MemOp, RoutingOp, is_monadic_alu


# Build mnemonic to opcode mapping
MNEMONIC_TO_OP: dict[str, Union[ArithOp, LogicOp, RoutingOp, MemOp]] = {
    # Arithmetic operations
    "add": ArithOp.ADD,
    "sub": ArithOp.SUB,
    "inc": ArithOp.INC,
    "dec": ArithOp.DEC,
    "shl": ArithOp.SHL,
    "shr": ArithOp.SHR,
    "asr": ArithOp.ASR,
    # Logic operations
    "and": LogicOp.AND,
    "or": LogicOp.OR,
    "xor": LogicOp.XOR,
    "not": LogicOp.NOT,
    "eq": LogicOp.EQ,
    "lt": LogicOp.LT,
    "lte": LogicOp.LTE,
    "gt": LogicOp.GT,
    "gte": LogicOp.GTE,
    # Routing/branch operations
    "breq": RoutingOp.BREQ,
    "brgt": RoutingOp.BRGT,
    "brge": RoutingOp.BRGE,
    "brof": RoutingOp.BROF,
    "sweq": RoutingOp.SWEQ,
    "swgt": RoutingOp.SWGT,
    "swge": RoutingOp.SWGE,
    "swof": RoutingOp.SWOF,
    "gate": RoutingOp.GATE,
    "sel": RoutingOp.SEL,
    "merge": RoutingOp.MRGE,
    "pass": RoutingOp.PASS,
    "const": RoutingOp.CONST,
    "free_frame": RoutingOp.FREE_FRAME,  # ALU free (deallocate frame slot)
    "extract_tag": RoutingOp.EXTRACT_TAG,
    "alloc_remote": RoutingOp.ALLOC_REMOTE,
    # Memory operations
    "read": MemOp.READ,
    "write": MemOp.WRITE,
    "clear": MemOp.CLEAR,
    "alloc": MemOp.ALLOC,
    "free": MemOp.FREE,  # SM free
    "rd_inc": MemOp.RD_INC,
    "rd_dec": MemOp.RD_DEC,
    "cmp_sw": MemOp.CMP_SW,
    "exec": MemOp.EXEC,
    "raw_read": MemOp.RAW_READ,
    "set_page": MemOp.SET_PAGE,
    "write_imm": MemOp.WRITE_IMM,
    "ext": MemOp.EXT,
}


# Build reverse mapping with type information to avoid IntEnum collisions
_reverse_mapping: dict[tuple[type, int], str] = {}
for mnemonic, op in MNEMONIC_TO_OP.items():
    _reverse_mapping[(type(op), int(op))] = mnemonic


class TypeAwareOpToMnemonicDict:
    """Collision-free reverse mapping from opcodes to mnemonics.

    Handles IntEnum cross-type equality by using (type, value) tuples internally.
    Supports dict-like access: OP_TO_MNEMONIC[ArithOp.ADD] returns "add",
    OP_TO_MNEMONIC[MemOp.READ] returns "read", etc.
    """

    def __init__(self, mapping: dict[tuple[type, int], str]):
        """Initialize with a type-indexed mapping.

        Args:
            mapping: dict from (type, value) tuples to mnemonic strings
        """
        self._mapping = mapping

    def __getitem__(self, op: Union[ArithOp, LogicOp, RoutingOp, MemOp]) -> str:
        """Get the mnemonic for an opcode.

        Args:
            op: The opcode enum value

        Returns:
            The mnemonic string

        Raises:
            KeyError: If the opcode is not in the mapping
        """
        key = (type(op), int(op))
        if key not in self._mapping:
            raise KeyError(f"Opcode {op} ({type(op).__name__}) not found in mapping")
        return self._mapping[key]

    def __contains__(self, op: Union[ArithOp, LogicOp, RoutingOp, MemOp]) -> bool:
        """Check if an opcode is in the mapping."""
        return (type(op), int(op)) in self._mapping

    def __iter__(self):
        """Iterate over mnemonic strings."""
        return iter(self._mapping.values())

    def __len__(self) -> int:
        """Return the number of opcode-mnemonic pairs."""
        return len(self._mapping)

    def items(self):
        """Return an iterator of (opcode_type, mnemonic) pairs for testing.

        This reconstructs the original enum instances from the stored types/values.
        """
        result = []
        for (op_type, op_val), mnemonic in self._mapping.items():
            result.append((op_type(op_val), mnemonic))
        return result


OP_TO_MNEMONIC: TypeAwareOpToMnemonicDict = TypeAwareOpToMnemonicDict(_reverse_mapping)


# Set of opcodes that are always monadic (single input operand).
# We use a frozenset of (type, value) tuples to avoid IntEnum collisions.
_MONADIC_OPS_TUPLES: frozenset[tuple[type, int]] = frozenset([
    # ALU ops: duplicated from cm_inst.is_monadic_alu() for MONADIC_OPS membership testing.
    # is_monadic() short-circuits to is_monadic_alu() for ALU ops, so these entries
    # are only reached via `op in MONADIC_OPS` (TypeAwareMonadicOpsSet).
    (ArithOp, int(ArithOp.INC)),
    (ArithOp, int(ArithOp.DEC)),
    (ArithOp, int(ArithOp.SHL)),
    (ArithOp, int(ArithOp.SHR)),
    (ArithOp, int(ArithOp.ASR)),
    # Logic: single input
    (LogicOp, int(LogicOp.NOT)),
    # Routing: single input or no ALU involvement
    (RoutingOp, int(RoutingOp.PASS)),
    (RoutingOp, int(RoutingOp.CONST)),
    (RoutingOp, int(RoutingOp.FREE_FRAME)),
    (RoutingOp, int(RoutingOp.EXTRACT_TAG)),
    (RoutingOp, int(RoutingOp.ALLOC_REMOTE)),
    # Memory: single input (monadic SM operations)
    (MemOp, int(MemOp.READ)),
    (MemOp, int(MemOp.ALLOC)),
    (MemOp, int(MemOp.FREE)),
    (MemOp, int(MemOp.CLEAR)),
    (MemOp, int(MemOp.RD_INC)),
    (MemOp, int(MemOp.RD_DEC)),
    (MemOp, int(MemOp.EXEC)),
    (MemOp, int(MemOp.RAW_READ)),
    (MemOp, int(MemOp.SET_PAGE)),
    (MemOp, int(MemOp.WRITE_IMM)),
    (MemOp, int(MemOp.EXT)),
])


class TypeAwareMonadicOpsSet:
    """Collision-free set of monadic opcodes.

    Handles IntEnum cross-type equality by using (type, value) tuples internally.
    Supports membership testing: ArithOp.INC in MONADIC_OPS returns True,
    but ArithOp.ADD in MONADIC_OPS returns False (collision-free).
    """

    def __init__(self, tuples: frozenset[tuple[type, int]]):
        """Initialize with type-indexed tuples.

        Args:
            tuples: frozenset of (type, value) tuples
        """
        self._tuples = tuples

    def __contains__(self, op: Union[ArithOp, LogicOp, RoutingOp, MemOp]) -> bool:
        """Check if an opcode is in the set, handling IntEnum collisions.

        Args:
            op: The opcode enum value

        Returns:
            True if the opcode is monadic, False otherwise
        """
        return (type(op), int(op)) in self._tuples

    def __iter__(self):
        """Iterate over opcode instances in the set."""
        for op_type, op_val in self._tuples:
            yield op_type(op_val)

    def __len__(self) -> int:
        """Return the number of monadic opcodes."""
        return len(self._tuples)

    def __repr__(self) -> str:
        """Return a string representation of the set."""
        ops = list(self)
        return f"TypeAwareMonadicOpsSet({ops})"


MONADIC_OPS: TypeAwareMonadicOpsSet = TypeAwareMonadicOpsSet(_MONADIC_OPS_TUPLES)


def is_monadic(op: Union[ArithOp, LogicOp, RoutingOp, MemOp], const: Optional[int] = None) -> bool:
    """Check if an opcode is monadic (single input operand).

    Args:
        op: The ALUOp or MemOp enum value
        const: Optional const value. Used to determine monadic form of WRITE.
               If const is not None, WRITE is monadic (cell_addr from const).
               If const is None, WRITE is dyadic (cell_addr from left operand).

    Returns:
        True if the opcode is always monadic, or if it's WRITE with const set.
        False for CMP_SW (always dyadic) and WRITE with const=None (dyadic).
    """
    # Use canonical is_monadic_alu for ALU operations
    if isinstance(op, (ArithOp, LogicOp, RoutingOp)):
        return is_monadic_alu(op)

    # Handle MemOp operations
    op_tuple = (type(op), int(op))
    if op_tuple in _MONADIC_OPS_TUPLES:
        return True

    # Special case: WRITE can be monadic (const given) or dyadic (const not given)
    if type(op) is MemOp and op == MemOp.WRITE:
        return const is not None

    return False


def is_dyadic(op: Union[ArithOp, LogicOp, RoutingOp, MemOp], const: Optional[int] = None) -> bool:
    """Check if an opcode is dyadic (two input operands).

    Args:
        op: The ALUOp or MemOp enum value
        const: Optional const value. Used for context-dependent operations like WRITE.

    Returns:
        True if the opcode is dyadic, False otherwise.
    """
    return not is_monadic(op, const)