import sys
import typing

from reprospect.test.sass.composite import instructions_contain
from reprospect.test.sass.composite_impl import SequenceMatcher
from reprospect.test.sass.instruction import (
    InstructionMatch,
    OpcodeModsWithOperandsMatcher,
    PatternBuilder,
    RegisterMatcher,
)
from reprospect.test.sass.instruction.constant import Constant
from reprospect.test.sass.instruction.instruction import ZeroOrOne
from reprospect.test.sass.instruction.register import Register
from reprospect.tools.sass.decode import Instruction

if sys.version_info >= (3, 12):
    from typing import override
else:
    from typing_extensions import override

class AddInt128Matcher(SequenceMatcher):
    """
    Match SASS instructions originating from the addition of 2 :code:`__int128`.

    It may use one of several instruction patterns:

    * :py:meth:`pattern_3IADD`
    * :py:meth:`pattern_4IADD3`

    .. note::

        For a simple "load, add, store" sequence, the PTX may look as follows::

            ld.global.nc.v2.u64 {%rd7, %rd8}, [%rd6];
            add.s64 %rd9, %rd3, %rd5;
            ld.global.v2.u64 {%rd10, %rd11}, [%rd9];
            add.cc.s64 %rd12, %rd10, %rd7;
            addc.cc.s64 %rd13, %rd11, %rd8;
            st.global.v2.u64 [%rd9], {%rd12, %rd13};

        According to https://docs.nvidia.com/cuda/parallel-thread-execution/#extended-precision-arithmetic-instructions-add-cc, the carry-out
        value is written in the condition code register, usually ``P0``.

    .. note::

        It is not decorated with :py:func:`dataclasses.dataclass` because of https://github.com/mypyc/mypyc/issues/1061.
    """
    __slots__ = ('_index', 'start')

    def __init__(self, start: str | None = None) -> None:
        self._index: int = 0
        self.start: typing.Final[str] = PatternBuilder.group(
            start or Register.mod(PatternBuilder.any(Register.REG, Register.UREG), reuse=None, captured=False), group='start',
        )

    def pattern_3IADD(self, instructions: typing.Sequence[Instruction | str]) -> list[InstructionMatch] | None: # pylint: disable=invalid-name
        """
        Typically::

            IADD.64 RZ, P0, R4, UR12
            IADD.64 R4, R4, UR12
            IADD.64.X R6, R6, UR14, P0
        """
        matched: list[InstructionMatch] = []

        offset = 0

        matcher_iadd_step_0 = OpcodeModsWithOperandsMatcher(
            opcode='IADD', modifiers=('64',),
            operands=(
                'RZ',
                Register.PRED,
                self.start,
                Register.mod(PatternBuilder.any(Register.REG, Register.UREG), reuse=None, captured=False),
            ),
        )
        if (matched_iadd_step_0 := matcher_iadd_step_0.match(instructions[offset])) is None:
            return None
        offset += 1
        matched.append(matched_iadd_step_0)

        assert (iadd_step_0_reg_0 := RegisterMatcher(special=False).match(matched_iadd_step_0.operands[-2])) is not None
        assert (iadd_step_0_reg_1 := RegisterMatcher(special=False).match(matched_iadd_step_0.operands[-1])) is not None

        matcher_iadd_step_1 = OpcodeModsWithOperandsMatcher(
            opcode='IADD', modifiers=('64',),
            operands=(
                Register.REG,
                f'{iadd_step_0_reg_0.rtype}{iadd_step_0_reg_0.index}',
                f'{iadd_step_0_reg_1.rtype}{iadd_step_0_reg_1.index}',
            ),
        )
        if (matched_iadd_step_1 := matcher_iadd_step_1.match(instructions[offset])) is None:
            return None
        offset += 1
        matched.append(matched_iadd_step_1)

        iadd_step_2_reg_0 = f'{iadd_step_0_reg_0.rtype}{iadd_step_0_reg_0.index + 2}'
        iadd_step_2_reg_1 = f'{iadd_step_0_reg_1.rtype}{iadd_step_0_reg_1.index + 2}'

        matcher_iadd_step_2 = instructions_contain(OpcodeModsWithOperandsMatcher(
            opcode='IADD', modifiers=('64', 'X'),
            operands=(
                Register.REG,
                iadd_step_2_reg_0,
                iadd_step_2_reg_1,
                matched_iadd_step_0.operands[1],
            ),
        ))
        if (matched_iadd_step_2 := matcher_iadd_step_2.match(instructions[offset:])) is None:
            return None
        self._index = offset + matcher_iadd_step_2.next_index

        matched.extend(matched_iadd_step_2)

        return matched

    def pattern_4IADD3(self, instructions: typing.Sequence[Instruction | str]) -> list[InstructionMatch] | None: # pylint: disable=invalid-name
        """
        Typically::

            IADD3 R20, P0, R8, c[0x0][0x180], RZ
            IADD3.X R21, P0, R9, c[0x0][0x184], RZ, P0, !PT
            IADD3.X R22, P0, R10, c[0x0][0x188], RZ, P0, !PT
            IADD3.X R23, R11, c[0x0][0x18c], RZ, P0, !PT

        or::

            IADD3 R4, P0, PT, R4, R12, RZ
            IADD3.X R5, P0, PT, R5, R13, RZ, P0, !PT
            IADD3.X R6, P0, PT, R6, R14, RZ, P0, !PT
            IADD3.X R7, PT, PT, R7, R15, RZ, P0, !PT

        That is, there may be one additional argument.
        """
        matched: list[InstructionMatch] = []

        offset = 0

        matcher_iadd3_step_0 = OpcodeModsWithOperandsMatcher(
            opcode='IADD3',
            operands=(
                Register.REG,
                Register.PRED,
                ZeroOrOne(Register.PREDT),
                self.start,
                PatternBuilder.any(Register.REG, Register.UREG, Constant.ADDRESS),
                'RZ',
            ),
        )
        if (matched_iadd3_step_0 := matcher_iadd3_step_0.match(instructions[offset])) is None:
            return None
        offset += 1
        matched.append(matched_iadd3_step_0)

        index_src = 2 if len(matched_iadd3_step_0.operands) == 5 else 3

        assert (iadd3_step_0_dst := RegisterMatcher(special=False).match(matched_iadd3_step_0.operands[0])) is not None
        assert (iadd3_step_0_src := RegisterMatcher(special=False).match(matched_iadd3_step_0.operands[index_src])) is not None
        iadd3_step_1_dst = f'{iadd3_step_0_dst.rtype}{iadd3_step_0_dst.index + 1}'
        iadd3_step_1_src = f'{iadd3_step_0_src.rtype}{iadd3_step_0_src.index + 1}'

        matcher_iadd3_step_1 = instructions_contain(OpcodeModsWithOperandsMatcher(
            opcode='IADD3', modifiers=('X',),
            operands=(
                iadd3_step_1_dst,
                Register.PRED,
                ZeroOrOne(Register.PREDT),
                iadd3_step_1_src,
                PatternBuilder.any(Register.REG, Register.UREG, Constant.ADDRESS),
                'RZ',
                Register.PRED,
                '!PT',
            ),
        ))
        if (matched_iadd3_step_1 := matcher_iadd3_step_1.match(instructions[offset:])) is None:
            return None
        offset += matcher_iadd3_step_1.next_index
        matched.extend(matched_iadd3_step_1)

        iadd3_step_2_dst = f'{iadd3_step_0_dst.rtype}{iadd3_step_0_dst.index + 2}'
        iadd3_step_2_src = f'{iadd3_step_0_src.rtype}{iadd3_step_0_src.index + 2}'

        matcher_iadd3_step_2 = instructions_contain(OpcodeModsWithOperandsMatcher(
            opcode='IADD3', modifiers=('X',),
            operands=(
                iadd3_step_2_dst,
                Register.PRED,
                ZeroOrOne(Register.PREDT),
                iadd3_step_2_src,
                PatternBuilder.any(Register.REG, Register.UREG, Constant.ADDRESS),
                'RZ',
                Register.PRED,
                '!PT',
            ),
        ))
        if (matched_iadd3_step_2 := matcher_iadd3_step_2.match(instructions[offset:])) is None:
            return None
        offset += matcher_iadd3_step_2.next_index
        matched.extend(matched_iadd3_step_2)

        iadd3_step_3_dst = f'{iadd3_step_0_dst.rtype}{iadd3_step_0_dst.index + 3}'
        iadd3_step_3_src = f'{iadd3_step_0_src.rtype}{iadd3_step_0_src.index + 3}'

        matcher_iadd3_step_3 = instructions_contain(OpcodeModsWithOperandsMatcher(
            opcode='IADD3', modifiers=('X',),
            operands=(
                iadd3_step_3_dst,
                ZeroOrOne(Register.PREDT),
                ZeroOrOne(Register.PREDT),
                iadd3_step_3_src,
                PatternBuilder.any(Register.REG, Register.UREG, Constant.ADDRESS),
                'RZ',
                Register.PRED,
                '!PT',
            ),
        ))
        if (matched_iadd3_step_3 := matcher_iadd3_step_3.match(instructions[offset:])) is None:
            return None
        self._index = offset + matcher_iadd3_step_3.next_index
        matched.extend(matched_iadd3_step_3)

        return matched

    @override
    def match(self, instructions: typing.Sequence[Instruction | str]) -> list[InstructionMatch] | None:
        instruction = instructions[0].instruction if isinstance(instructions[0], Instruction) else instructions[0]
        if instruction.startswith('IADD.64'):
            return self.pattern_3IADD(instructions=instructions)
        if instruction.startswith('IADD3'):
            return self.pattern_4IADD3(instructions=instructions)
        return None

    @override
    @property
    def next_index(self) -> int:
        return self._index
