from typing import Optional

from tensorrt_llm.functional import Tensor, silu
from tensorrt_llm.layers import ColumnLinear
from tensorrt_llm.mapping import Mapping
from tensorrt_llm.module import Module, ModuleList

from ..._utils import str_dtype_to_trt


class ResBlock(Module):

    def __init__(self,
                 exit_dim: int,
                 dtype: Optional[str],
                 mapping: Mapping = Mapping()):
        super().__init__()
        self.linear = ColumnLinear(
            exit_dim,
            exit_dim,
            bias=True,
            dtype=dtype,
            tp_group=mapping.tp_group,
            tp_size=mapping.tp_size,
            gather_output=True,
        )

    def forward(self, x: Tensor) -> Tensor:
        return x + silu(self.linear(x))


class Drafter(Module):

    def __init__(
            self,
            num_layers: int,
            hidden_size: int,
            exit_dim: int,
            vocab_size: int,
            dtype: Optional[str] = None,
            is_rnn: bool = False,
            mapping: Mapping = Mapping(),
    ):
        super().__init__()
        self.num_layers = num_layers
        self.is_rnn = is_rnn
        self.dtype = str_dtype_to_trt(dtype)

        input_dim = 2 * hidden_size
        self.input_proj = (None if input_dim == exit_dim else ColumnLinear(
            input_dim,
            exit_dim,
            bias=True,
            dtype=dtype,
            tp_group=mapping.tp_group,
            tp_size=mapping.tp_size,
            gather_output=True,
        ))

        self.layers = ModuleList([
            ResBlock(exit_dim, dtype, mapping) for _ in range(self.num_layers)
        ])
        self.lm_head = ColumnLinear(
            exit_dim,
            vocab_size,
            bias=False,
            dtype=dtype,
            tp_group=mapping.tp_group,
            tp_size=mapping.tp_size,
            gather_output=True,
        )

        if is_rnn:
            self.rnn_u = ColumnLinear(
                hidden_size,
                hidden_size,
                bias=True,
                dtype=dtype,
                tp_group=mapping.tp_group,
                tp_size=mapping.tp_size,
                gather_output=True,
            )
            self.rnn_w = ColumnLinear(
                hidden_size,
                hidden_size,
                bias=False,
                dtype=dtype,
                tp_group=mapping.tp_group,
                tp_size=mapping.tp_size,
                gather_output=True,
            )
        return

    @classmethod
    def from_config(cls, config, vocab_size_padded):
        kwargs = {
            "num_layers": config.redrafter_num_layers,
            "hidden_size": config.redrafter_hidden_size,
            "exit_dim": config.redrafter_exit_dim,
            "vocab_size": vocab_size_padded,
            "dtype": config.dtype,
            "is_rnn": config.redrafter_is_rnn,
            "mapping": config.mapping,
        }
        return cls(**kwargs)

    def forward(self, x: Tensor) -> Tensor:
        hidden_states = self.input_proj(x) if self.input_proj is not None else x
        for layer in self.layers:
            hidden_states = layer(hidden_states)

        return self.lm_head(hidden_states)

    def rnn_embed(self, x: Tensor, prev: Tensor = None) -> Tensor:
        assert self.is_rnn, "This function should not be called when redrafter_is_rnn is false."
        w_embd = self.rnn_w(x)
        return w_embd if prev is None else w_embd + self.rnn_u(prev)