"""Opcode-to-category mapping for visual graph rendering.

Maps each ALUOp/MemOp to a visual category and colour
for the dataflow graph renderer.
"""

from __future__ import annotations

from enum import Enum
from typing import Union

from cm_inst import ArithOp, LogicOp, MemOp, RoutingOp


class OpcodeCategory(Enum):
    ARITHMETIC = "arithmetic"
    LOGIC = "logic"
    COMPARISON = "comparison"
    ROUTING = "routing"
    MEMORY = "memory"
    IO = "io"  # reserved for future I/O ops (ior, iow, iorw) — not yet in asm/opcodes.py
    CONFIG = "config"
    STRUCTURE_MEMORY = "structure_memory"


CATEGORY_COLOURS: dict[OpcodeCategory, str] = {
    OpcodeCategory.ARITHMETIC: "#4a90d9",
    OpcodeCategory.LOGIC: "#4caf50",
    OpcodeCategory.COMPARISON: "#ff9800",
    OpcodeCategory.ROUTING: "#9c27b0",
    OpcodeCategory.MEMORY: "#ff5722",
    OpcodeCategory.IO: "#009688",
    OpcodeCategory.CONFIG: "#9e9e9e",
    OpcodeCategory.STRUCTURE_MEMORY: "#795548",
}


_COMPARISON_OPS: frozenset[LogicOp] = frozenset({
    LogicOp.EQ, LogicOp.LT, LogicOp.LTE, LogicOp.GT, LogicOp.GTE,
})

_CONFIG_ROUTING_OPS: frozenset[RoutingOp] = frozenset({
    RoutingOp.CONST, RoutingOp.FREE_FRAME,
})


def categorise(op: Union[ArithOp, LogicOp, RoutingOp, MemOp]) -> OpcodeCategory:
    """Categorise an opcode for visual rendering.

    Maps each opcode to a visual category used by the graph renderer.
    Handles special cases like LogicOp comparison ops and RoutingOp config ops.

    Args:
        op: An opcode enum value (ArithOp, LogicOp, RoutingOp, or MemOp)

    Returns:
        The OpcodeCategory for this opcode

    Raises:
        ValueError: If the opcode type is unknown
    """
    if isinstance(op, ArithOp):
        return OpcodeCategory.ARITHMETIC
    if isinstance(op, LogicOp):
        if op in _COMPARISON_OPS:
            return OpcodeCategory.COMPARISON
        return OpcodeCategory.LOGIC
    if isinstance(op, RoutingOp):
        if op in _CONFIG_ROUTING_OPS:
            return OpcodeCategory.CONFIG
        return OpcodeCategory.ROUTING
    if isinstance(op, MemOp):
        return OpcodeCategory.MEMORY
    raise ValueError(f"Unknown opcode type: {type(op).__name__}")