import ida_hexrays

from d810.core import getLogger
from d810.hexrays.hexrays_helpers import append_mop_if_not_in_list
from d810.optimizers.microcode.handler import ConfigParam
from d810.optimizers.microcode.flow.flattening.generic import (
    GenericDispatcherBlockInfo,
    GenericDispatcherCollector,
    GenericDispatcherInfo,
    GenericDispatcherUnflatteningRule,
)

unflat_logger = getLogger("D810.unflat")
FLATTENING_JUMP_OPCODES = [ida_hexrays.m_jtbl]


class TigressSwitchDispatcherBlockInfo(GenericDispatcherBlockInfo):
    pass


class TigressSwitchDispatcherInfo(GenericDispatcherInfo):
    def explore(self, blk: ida_hexrays.mblock_t):
        self.reset()
        if not self._is_candidate_for_dispatcher_entry_block(blk):
            return False
        self.mop_compared, mcases = self._get_comparison_info(blk)
        assert mcases is not None
        self.entry_block = TigressSwitchDispatcherBlockInfo(blk)
        self.entry_block.parse()
        for used_mop in self.entry_block.use_list:
            append_mop_if_not_in_list(used_mop, self.entry_block.assume_def_list)
        self.dispatcher_internal_blocks.append(self.entry_block)
        for possible_values, target_block_serial in zip(
            mcases.c.values, mcases.c.targets
        ):
            if target_block_serial == self.entry_block.serial:
                continue
            exit_block = TigressSwitchDispatcherBlockInfo(
                blk.mba.get_mblock(target_block_serial), self.entry_block
            )
            self.dispatcher_exit_blocks.append(exit_block)
            if len(possible_values) == 0:
                continue
            self.comparison_values.append(possible_values[0])
        return True

    def _get_comparison_info(self, blk: ida_hexrays.mblock_t):
        # blk.tail must be a jtbl
        if (blk.tail is None) or (blk.tail.opcode != ida_hexrays.m_jtbl):
            return None, None
        return blk.tail.l, blk.tail.r

    def _is_candidate_for_dispatcher_entry_block(self, blk: ida_hexrays.mblock_t):
        if (blk.tail is None) or (blk.tail.opcode != ida_hexrays.m_jtbl):
            return False
        return True


class TigressSwitchDispatcherCollector(GenericDispatcherCollector):
    DISPATCHER_CLASS = TigressSwitchDispatcherInfo
    DEFAULT_DISPATCHER_MIN_INTERNAL_BLOCK = 0
    DEFAULT_DISPATCHER_MIN_EXIT_BLOCK = 2
    DEFAULT_DISPATCHER_MIN_COMPARISON_VALUE = 2


class UnflattenerSwitchCase(GenericDispatcherUnflatteningRule):
    DESCRIPTION = "Remove control flow flattening generated by Tigress with Switch case dispatcher"
    CONFIG_SCHEMA = GenericDispatcherUnflatteningRule.CONFIG_SCHEMA + (
        ConfigParam(
            "max_exit_blocks_for_additional_passes",
            int,
            256,
            "For huge switch dispatchers, skip pass>0 to avoid unstable CFG rewrite blowups",
        ),
    )
    # MMAT_GLBOPT3 is rarely/never called by Hex-Rays
    DEFAULT_UNFLATTENING_MATURITIES = [
        ida_hexrays.MMAT_CALLS,
        ida_hexrays.MMAT_GLBOPT1,
        ida_hexrays.MMAT_GLBOPT2,
    ]
    DEFAULT_MAX_DUPLICATION_PASSES = 20
    DEFAULT_MAX_PASSES = 7
    DEFAULT_MAX_EXIT_BLOCKS_FOR_ADDITIONAL_PASSES = 256

    def __init__(self):
        super().__init__()
        self.max_exit_blocks_for_additional_passes = (
            self.DEFAULT_MAX_EXIT_BLOCKS_FOR_ADDITIONAL_PASSES
        )
        # Large switch dispatchers can crash Hex-Rays when rewritten repeatedly
        # across maturities. Track functions that already consumed one pass.
        self._large_dispatcher_single_pass_done: set[int] = set()

    def configure(self, kwargs):
        super().configure(kwargs)
        if "max_exit_blocks_for_additional_passes" in self.config:
            self.max_exit_blocks_for_additional_passes = int(
                self.config["max_exit_blocks_for_additional_passes"]
            )

    def should_skip_pass_for_layout(self, layout_signals):
        threshold = int(self.max_exit_blocks_for_additional_passes)
        if threshold <= 0:
            return False, None
        max_exit_blocks = int(layout_signals.get("max_exit_blocks", 0))
        if max_exit_blocks <= threshold:
            return False, None
        func_ea = int(getattr(self.mba, "entry_ea", 0))
        if func_ea in self._large_dispatcher_single_pass_done:
            return (
                True,
                "UnflattenerSwitchCase: skipping additional pass at maturity %s "
                "for large dispatcher in function 0x%x "
                "(max_exit_blocks=%d > threshold=%d)"
                % (
                    self.mba.maturity,
                    func_ea,
                    max_exit_blocks,
                    threshold,
                ),
            )
        self._large_dispatcher_single_pass_done.add(func_ea)
        # Allow exactly one pass for this large dispatcher (first seen pass).
        if self.cur_maturity_pass <= 0:
            return False, None
        return (
            True,
            "UnflattenerSwitchCase: skipping pass %d at maturity %s for large "
            "dispatcher (max_exit_blocks=%d > threshold=%d)"
            % (
                self.cur_maturity_pass,
                self.mba.maturity,
                max_exit_blocks,
                threshold,
            ),
        )

    @property
    def DISPATCHER_COLLECTOR_CLASS(self) -> type[GenericDispatcherCollector]:
        """Return the class of the dispatcher collector."""
        return TigressSwitchDispatcherCollector
