"""
This dialect provides operations to target features of the Snitch[1]
streaming architecture based on custom extensions to the RISC-V ISA.
This dialect works on 'riscv' types directly as all arguments are of
'riscv.reg<>' type and it is meant to be as close as possible to the asm
that aims at generating.

[1] https://pulp-platform.github.io/snitch/publications
"""

from abc import ABC
from collections.abc import Sequence
from dataclasses import dataclass

from xdsl.dialects import stream
from xdsl.dialects.builtin import IntAttr
from xdsl.dialects.riscv import IntRegisterType
from xdsl.ir import Attribute, Dialect, Operation, SSAValue
from xdsl.irdl import (
    IRDLOperation,
    Operand,
    attr_def,
    irdl_op_definition,
    operand_def,
    var_result_def,
)
from xdsl.utils.exceptions import VerifyException


@dataclass(frozen=True)
class SnitchResources:
    """
    Bounds for resources provided by the Snitch architecture.
    """

    # Number of dimensions supported by each data mover.
    dimensions: int = 4


class SsrSetDimensionConfigOperation(IRDLOperation, ABC):
    """
    A base class for Snitch operations that set a
    configuration value for a specific dimension handled by a streamer.
    """

    value: Operand = operand_def(IntRegisterType)
    dm = attr_def(IntAttr)
    dimension = attr_def(IntAttr)

    def __init__(
        self,
        value: Operation | SSAValue,
        dm: IntAttr,
        dimension: IntAttr,
    ):
        super().__init__(
            operands=[value],
            attributes={
                "dm": dm,
                "dimension": dimension,
            },
        )

    def verify_(self) -> None:
        if self.dimension.data >= SnitchResources.dimensions:
            raise VerifyException(
                f"dimension attribute out of range [0..{SnitchResources.dimensions-1}], "
                f"Snitch supports up to {SnitchResources.dimensions} dimensions per streamer"
            )


class SsrSetStreamConfigOperation(IRDLOperation, ABC):
    """
    A base class for Snitch operations that set a
    configuration value for a streamer.
    """

    value: Operand = operand_def(IntRegisterType)
    dm = attr_def(IntAttr)

    def __init__(self, value: Operation | SSAValue, dm: IntAttr):
        super().__init__(
            operands=[value],
            attributes={
                "dm": dm,
            },
        )


@irdl_op_definition
class SsrSetDimensionBoundOp(SsrSetDimensionConfigOperation):
    """
    Set the bound for one of the dimensions handled by a
    specific streamer.
    """

    name = "snitch.ssr_set_dimension_bound"


@irdl_op_definition
class SsrSetDimensionStrideOp(SsrSetDimensionConfigOperation):
    """
    Set the stride for one of the dimensions handled by a
    specific streamer.
    """

    name = "snitch.ssr_set_dimension_stride"


@irdl_op_definition
class SsrSetDimensionSourceOp(SsrSetDimensionConfigOperation):
    """
    Set the data source for one of the dimensions handled by a
    specific streamer.
    """

    name = "snitch.ssr_set_dimension_source"


@irdl_op_definition
class SsrSetDimensionDestinationOp(SsrSetDimensionConfigOperation):
    """
    Set the data destination for one of the dimensions handled by a
    specific streamer.
    """

    name = "snitch.ssr_set_dimension_destination"


@irdl_op_definition
class SsrSetStreamRepetitionOp(SsrSetStreamConfigOperation):
    """
    Setup repetition count for a specific data mover.
    """

    name = "snitch.ssr_set_stream_repetition"


@irdl_op_definition
class SsrEnable(IRDLOperation):
    """
    Enable stream semantics.
    """

    name = "snitch.ssr_enable"

    streams = var_result_def(stream.StreamType)

    def __init__(self, stream_types: Sequence[Attribute]):
        super().__init__(result_types=[stream_types])


@irdl_op_definition
class SsrDisable(IRDLOperation):
    """
    Disable stream semantics.
    """

    name = "snitch.ssr_disable"

    def __init__(self):
        super().__init__()


Snitch = Dialect(
    "snitch",
    [
        SsrSetDimensionBoundOp,
        SsrSetDimensionStrideOp,
        SsrSetDimensionSourceOp,
        SsrSetDimensionDestinationOp,
        SsrSetStreamRepetitionOp,
        SsrEnable,
        SsrDisable,
    ],
    [],
)
