import dataclasses
import io
import logging
import pathlib
import re
import sys
import typing

import mypy_extensions
import pandas
import rich.table

from reprospect.utils import rich_helpers

if sys.version_info >= (3, 11):
    from enum import StrEnum
else:
    from backports.strenum.strenum import StrEnum

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

@dataclasses.dataclass(frozen = True, slots = True)
class ControlCode:
    """
    SASS control code decoding.

    SASS instructions use 128-bit encoding:

    * 64-bit instruction word
    * 64-bit word including control code

    The control code manages instruction dependencies, scheduling, and resource usage.

    The control code is a section of 21 bits with the following structure:

    +------------------+---------------------------------------------+
    | number of bits   | meaning                                     |
    +==================+=============================================+
    | 4                | stall count (0-15 cycles)                   |
    +------------------+---------------------------------------------+
    | 1                | yield flag                                  |
    +------------------+---------------------------------------------+
    | 3                | write barrier                               |
    +------------------+---------------------------------------------+
    | 3                | read barrier                                |
    +------------------+---------------------------------------------+
    | 6                | wait barrier mask (one bit per barrier 0-5) |
    +------------------+---------------------------------------------+
    | 4                | reuse flags mask (for operands A, B, C, D)  |
    +------------------+---------------------------------------------+

    Barriers
    ========

    Coordinate dependencies between instructions executing on different functional units:

    * Write barrier is set when the instruction completes and produces a result.
      It marks that the register is ready. Later instructions reading that register must wait.
    * Read barrier indicates that the instruction needs a register with the result from a previous instruction
      to remain unchanged. Later instructions writing on that register must wait.
    * The wait barrier mask is a bit mask with the barriers that the instruction must wait for
      (*e.g.* write and read barriers set by preceding instructions).

    A write or read barrier value with the three bits set to 1, i.e., an integer value of 7, means that no such barrier is set.

    Note that barriers are related to [scoreboarding](https://en.wikipedia.org/wiki/Scoreboarding).

    The reader may refer to the following references for further description of barriers and scoreboard dependencies:

    * :cite:`jia-2018-dissecting-nvidia-volta-gpu` (section 2.1)
    * :cite:`huerta-2025-analyzing-modern-nvidia-gpu` (figure 2).

    Stall count
    ===========

    Explicitly insert delay cycles to account for latency.

    It can be set within the 0-15 cycles range.

    Reuse flags
    ===========

    Operand forwarding optimization - reuse data from previous instruction.
    There are 4 reuse flags (`A`, `B`, `C` and `D`).

    If a subsequent instruction uses the same register in the same operand slot,
    for instance operand slot `B`, the compiler may set `reuse_B` to avoid register traffic.

    Yield flag
    ==========

    Hint to scheduler to yield to other warps. It is mainly used for long-latency operations and
    allows better multi-warp interleaving, thus improving throughput.

    Notes
    =====

    * There are 6 barriers (6 scoreboard barriers for dependency tracking).
    * Each warp has its own set of barriers.
    * The GPU hardware automatically enforces barrier dependencies.
    * The control code is generated by the CUDA ``ptxas`` tool.

    References:

    * :cite:`maxas-wiki-control-codes`
    """
    stall_count: int       #: Stall count.
    yield_flag : bool      #: Yield flag.
    read : int             #: Index of the read barrier that is set (7 if no read barrier is set).
    write: int             #: Index of the write barrier that is set (7 if no write barrier is set).
    wait : list[bool]      #: Wait barrier mask (one bit per barrier 0-5).
    reuse: dict[str, bool] #: Reuse flags.

    @staticmethod
    def decode(*, code : str) -> 'ControlCode':
        """
        Decode 64-bit word including a control code.
        """
        # The 41 first bits encode something else.
        control = int(code, 16) >> 41

        # 4 bits stall count.
        stall_count = control & 0xF

        # 1 bit yield flag.
        yield_flag = bool((control >> 4) & 0x1)

        # 3 bits write.
        write_barrier = (control >> 5) & 0x7

        # 3 bits read.
        read_barrier = (control >> 8) & 0x7

        # 6 bits wait barrier mask.
        wait_mask = (control >> 11) & 0x3F
        wait_barriers = [bool((wait_mask >> i) & 1) for i in range(6)]

        reuse_flags = {
            'A': bool((control >> 17) & 1),
            'B': bool((control >> 18) & 1),
            'C': bool((control >> 19) & 1),
            'D': bool((control >> 20) & 1)
        }

        return ControlCode(
            stall_count = stall_count,
            yield_flag = yield_flag,
            read = read_barrier,
            write = write_barrier,
            wait = wait_barriers,
            reuse = reuse_flags,
        )

class RegisterType(StrEnum):
    """
    Register types:

    * ``GPR``: General Purpose Registers.
    * ``PRED``: Predicate Registers.
    * ``UGPR``: Uniform General Purpose Registers.
    * ``UPRED``: Uniform Predicate Registers.
    """
    GPR = 'R'
    PRED = 'P'
    UGPR = 'UR'
    UPRED = 'UP'

@mypy_extensions.mypyc_attr(native_class = True)
@dataclasses.dataclass(frozen = True, slots = True)
class Instruction:
    """
    Represents a single SASS instruction with its components.
    """
    offset : int          #: Offset of the instruction in the SASS code.
    instruction : str     #: The disassembled SASS instruction including opcode, modifiers and operands.
    hex : str             #: The hexadecimal representation of the instruction.
    control : ControlCode #: The decoded control code associated with the instruction.

@mypy_extensions.mypyc_attr(native_class = True)
class Decoder(rich_helpers.TableMixin):
    """
    Parse the SASS assembly code extracted from a binary.

    The disassembled :py:attr:`instructions` are collected, and the associated control codes are decoded.
    """
    OFFSET : typing.Final[str] = r'[a-f0-9]+'

    HEX : typing.Final[str] = r'0x[a-f0-9]+'
    """
    Matcher for an hex-like string, such as `0x00000a0000017a02`.
    """

    MATCHER_CONTROL : typing.Final[re.Pattern[str]] = re.compile(rf'\/\* ({HEX}) \*\/')

    MATCHER : typing.Final[re.Pattern[str]] = re.compile(
        rf'/\*({OFFSET})\*/'
        r'\s+'
        r'(.*?)(?=\s{2,}|[;?&])'
        r'.*?'
        rf'/\*\s*({HEX})\s*\*/'
    )
    """
    Matcher for the full SASS line. It focuses on the offset, instruction and trailing hex encoding.

    For instance, it matches::

        /*0070*/                   UIMAD UR4, UR4, UR5, URZ ;                       /* 0x00000005040472a4 */

    Note that sometimes, additional strings appear after the instruction and before the hex::

        /*0090*/ ISETP.GE.U32.AND P0, PT, R0, UR9, PT            ?WAIT13_END_GROUP;  /* 0x0000000900007c0c */
        /*00e0*/ LDC.64 R10, c[0x0][0x3a0]  &wr=0x2  ?trans8;    /* 0x0000e800ff0a7b82 */

    These strings are ignored.
    """

    def __init__(self, source : typing.Optional[pathlib.Path] = None, code : typing.Optional[str] = None, skip_until_headerflags : bool = True) -> None:
        """
        Initialize the decoder with the SASS contained in `source` or `code`.
        """
        self.source = source
        self.code   = code
        self.instructions: list[Instruction] = [] #: The parsed instructions.
        if source or code:
            self._parse(skip_until_headerflags = skip_until_headerflags)

    def _parse(self, skip_until_headerflags : bool) -> None:
        """
        Parse SASS lines.
        """
        if self.source is not None:
            with self.source.open('r', encoding = 'utf-8') as fin:
                return self._parse_lines(fin, skip_until_headerflags)

        if self.code is not None:
            with io.StringIO(self.code) as fin:
                return self._parse_lines(fin, skip_until_headerflags)

        raise RuntimeError('Neither code nor source was given.')

    def _parse_lines(self, lines : typing.Iterable[str], skip_until_headerflags : bool) -> None:
        headerflags = False

        iterator : typing.Iterator[str] = iter(lines)

        for line in iterator:
            line = line.strip()

            # The line containing '..........' means the end of the SASS code.
            if skip_until_headerflags:
                if '..........' in line:
                    break

            # Skip empty lines.
            if not line:
                continue

            # Skip lines until '.headerflags' is met.
            if skip_until_headerflags and not headerflags:
                headerflags = '.headerflags' in line
                continue

            matchl = self.MATCHER.match(line)

            if not matchl:
                logging.error(f'The line:\n\t{line}\ndid not match {self.MATCHER}.')
                raise RuntimeError(line)

            # Peek next line safely (lookahead).
            try:
                next_line = next(iterator)
            except StopIteration:
                next_line = None

            control : ControlCode | None = None

            if next_line and (matchc := self.MATCHER_CONTROL.search(next_line)):
                control = ControlCode.decode(code = matchc.group(1))

            assert control is not None

            # Create instruction.
            self.instructions.append(Instruction(
                offset = int(matchl.group(1), base = 16),
                instruction = matchl.group(2).rstrip(),
                hex = matchl.group(3),
                control = control,
            ))

    def to_df(self) -> pandas.DataFrame:
        """
        Convert the decoded SASS to a :py:class:`pandas.DataFrame`.
        """
        data : dict[str, list[int | str]] = {
            'offset' : [],
            'instruction' : [],
            'stall' : [],
            'yield' : [],
            'b0' : [],
            'b1' : [],
            'b2' : [],
            'b3' : [],
            'b4' : [],
            'b5' : [],
        }

        for instruction in self.instructions:
            data['offset'].append(f'{instruction.offset:04x}')
            data['instruction'].append(instruction.instruction)
            data['stall'].append(instruction.control.stall_count)
            data['yield'].append(instruction.control.yield_flag)
            for i in range(6):
                value = []
                if instruction.control.wait[i]:
                    value.append('Wa')
                if instruction.control.write == i:
                    value.append('Wr')
                if instruction.control.read == i:
                    value.append('Re')
                data[f'b{i}'].append('/'.join(value) if value else '')

        return pandas.DataFrame(data)

    @override
    def to_table(self) -> rich.table.Table:
        """
        Convert the decoded SASS to a :py:class:`rich.table.Table`.
        """
        return rich_helpers.df_to_table(self.to_df())

    def to_html(self) -> str:
        """
        Visualize the decoded SASS in a nice HTML table.
        """
        return self.to_df().to_html()
