import unittest
from copy import deepcopy
from dataclasses import dataclass

import torch
from parameterized import parameterized
from transformers import Gemma3Config
from transformers import Gemma3ForCausalLM as HFGemma3ForCausalLM
from transformers.cache_utils import HybridCache

import tensorrt_llm
from tensorrt_llm._torch.attention_backend.utils import get_attention_backend
from tensorrt_llm._torch.metadata import KVCacheParams
from tensorrt_llm._torch.model_config import ModelConfig
from tensorrt_llm._torch.models.modeling_gemma3 import Gemma3ForCausalLM
from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
from tensorrt_llm.bindings.executor import KvCacheConfig
from tensorrt_llm.mapping import Mapping

# This is copied from https://huggingface.co/google/gemma-3-1b-it/blob/main/config.json.
# Updated to have 1 local layer and 1 global layer. Sliding window size updated to 4.
GEMMA3_1B_MINI_CONFIG = {
    "architectures": ["Gemma3ForCausalLM"],
    "attention_bias": False,
    "attention_dropout": 0.0,
    "attn_logit_softcapping": None,
    "bos_token_id": 2,
    "cache_implementation": "hybrid",
    "eos_token_id": [1, 106],
    "final_logit_softcapping": None,
    "head_dim": 256,
    "hidden_activation": "gelu_pytorch_tanh",
    "hidden_size": 1152,
    "initializer_range": 0.02,
    "intermediate_size": 6912,
    "max_position_embeddings": 32768,
    "model_type": "gemma3_text",
    "num_attention_heads": 4,
    "num_hidden_layers": 2,  # Modified for testing.
    "num_key_value_heads": 1,
    "pad_token_id": 0,
    "query_pre_attn_scalar": 256,
    "rms_norm_eps": 1e-06,
    "rope_local_base_freq": 10000,
    "rope_scaling": None,
    "rope_theta": 1000000,
    "sliding_window": 4,  # Modified for testing.
    "sliding_window_pattern": 2,  # Modified for testing.
    "torch_dtype": "bfloat16",
    "transformers_version": "4.50.0.dev0",
    "use_cache": True,
    "vocab_size": 262144
}


@dataclass(repr=False)
class Scenario:
    backend: str

    def __repr__(self) -> str:
        return f"backend:{self.backend.lower()}"


class TestGemma3(unittest.TestCase):

    def get_kv_cache_manager(self, dtype: torch.dtype, config: Gemma3Config,
                             tokens_per_block: int, max_seq_len: int,
                             batch_size: int, num_blocks: int):
        if dtype == torch.half:
            kv_cache_dtype = tensorrt_llm.bindings.DataType.HALF
        elif dtype == torch.bfloat16:
            kv_cache_dtype = tensorrt_llm.bindings.DataType.BF16
        else:
            raise ValueError("Invalid dtype")

        mapping = Mapping(world_size=1, tp_size=1, rank=0)
        kv_cache_config = KvCacheConfig(enable_block_reuse=False,
                                        enable_partial_reuse=False,
                                        copy_on_partial_reuse=False,
                                        max_tokens=num_blocks *
                                        tokens_per_block)
        kv_cache_manager = KVCacheManager(
            kv_cache_config,
            tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF,
            num_layers=config.num_hidden_layers,
            num_kv_heads=config.num_key_value_heads,
            head_dim=config.head_dim,
            tokens_per_block=tokens_per_block,
            max_seq_len=max_seq_len,
            max_batch_size=batch_size,
            mapping=mapping,
            dtype=kv_cache_dtype,
        )
        return kv_cache_manager

    def test_gemma3_sanity(self):

        config_dict = deepcopy(GEMMA3_1B_MINI_CONFIG)
        gemma3_config = Gemma3Config.from_dict(config_dict)

        dtype = gemma3_config.torch_dtype
        device = torch.device('cuda')

        model_config = ModelConfig(pretrained_config=gemma3_config)
        gemma3 = Gemma3ForCausalLM(model_config).to(device)

        input_ids = torch.tensor([100, 200, 300, 400, 500, 600, 700, 800],
                                 dtype=torch.int,
                                 device=device)

        context_sequence_lengths = [3, 2, 1]
        sequence_lengths = context_sequence_lengths + [1, 1]
        past_seen_tokens = [0, 0, 0, 62, 75]
        request_ids = list(range(len(sequence_lengths)))
        token_nums = (torch.tensor(past_seen_tokens) +
                      torch.tensor(sequence_lengths)).tolist()
        prompt_lens = token_nums[:3] + past_seen_tokens[3:]

        num_blocks = 100
        tokens_per_block = 128
        max_seq_len = num_blocks * tokens_per_block
        batch_size = len(context_sequence_lengths) + 2
        kv_cache_manager = self.get_kv_cache_manager(
            dtype=dtype,
            config=gemma3_config,
            tokens_per_block=tokens_per_block,
            max_seq_len=max_seq_len,
            batch_size=batch_size,
            num_blocks=num_blocks)
        kv_cache_manager.add_dummy_requests(request_ids, token_nums)

        metadata_cls = get_attention_backend(model_config.attn_backend).Metadata
        attn_metadata = metadata_cls(
            seq_lens=torch.tensor(sequence_lengths, dtype=torch.int),
            num_contexts=len(context_sequence_lengths),
            kv_cache_params=KVCacheParams(
                use_cache=True,
                num_cached_tokens_per_seq=past_seen_tokens,
            ),
            kv_cache_manager=kv_cache_manager,
            request_ids=request_ids,
            prompt_lens=prompt_lens,
            max_num_requests=len(context_sequence_lengths) + 2,
            max_num_tokens=8192,
        )

        position_ids = []
        for i, tokens in enumerate(past_seen_tokens):
            seq_len = context_sequence_lengths[i] if i < len(
                context_sequence_lengths) else 1
            position_id = torch.arange(tokens,
                                       tokens + seq_len,
                                       device=input_ids.device)
            position_ids.append(position_id)

        position_ids = torch.cat(position_ids).unsqueeze(0)

        with torch.inference_mode():
            attn_metadata.prepare()
            logits = gemma3.forward(input_ids=input_ids,
                                    position_ids=position_ids,
                                    attn_metadata=attn_metadata)

        self.assertEqual(len(past_seen_tokens), logits.shape[0])

        with torch.inference_mode():
            attn_metadata.prepare()
            logits = gemma3.forward(input_ids=input_ids,
                                    position_ids=position_ids,
                                    attn_metadata=attn_metadata,
                                    return_context_logits=True)
        self.assertEqual(input_ids.shape, logits.shape[:-1])

        kv_cache_manager.shutdown()

    @parameterized.expand([
        Scenario(backend="TRTLLM"),
        Scenario(backend="VANILLA"),
    ], lambda testcase_func, param_num, param:
                          f"{testcase_func.__name__}[{param.args[0]}]")
    @torch.no_grad()
    def test_gemma3_allclose_to_hf(self, scenario: Scenario) -> None:
        """
        Compare output to HF.
        """
        backend = scenario.backend
        metadata_cls = get_attention_backend(backend).Metadata

        torch.random.manual_seed(0)
        config_dict = deepcopy(GEMMA3_1B_MINI_CONFIG)
        gemma3_config = Gemma3Config.from_dict(config_dict)
        dtype = gemma3_config.torch_dtype
        device = torch.device('cuda')

        num_blocks = 1
        tokens_per_block = 128
        max_seq_len = num_blocks * tokens_per_block
        batch_size = 1

        hf_gemma3 = HFGemma3ForCausalLM(gemma3_config).to(dtype).to(
            device).eval()
        hf_cache = HybridCache(config=gemma3_config,
                               max_batch_size=batch_size,
                               max_cache_len=10,
                               device=device,
                               dtype=dtype)

        model_config = ModelConfig(pretrained_config=gemma3_config,
                                   attn_backend=backend)
        gemma3 = Gemma3ForCausalLM(model_config).to(dtype).to(device)
        gemma3.load_weights(hf_gemma3.state_dict())

        kv_cache_manager = self.get_kv_cache_manager(
            dtype=dtype,
            config=gemma3_config,
            tokens_per_block=tokens_per_block,
            max_seq_len=max_seq_len,
            batch_size=batch_size,
            num_blocks=num_blocks)

        # Context phase.
        input_ids = torch.tensor([100, 200, 300, 400, 500, 600, 700, 800],
                                 dtype=torch.int32,
                                 device=device)
        num_cached_tokens_per_seq = [0]
        request_ids = [1]
        token_nums = [input_ids.size(-1)]
        prompt_lens = [input_ids.size(-1)]
        kv_cache_manager.add_dummy_requests(request_ids, token_nums)

        attn_metadata = metadata_cls(
            seq_lens=torch.tensor([input_ids.size(-1)], dtype=torch.int),
            num_contexts=1,
            kv_cache_params=KVCacheParams(
                use_cache=True,
                num_cached_tokens_per_seq=num_cached_tokens_per_seq,
            ),
            max_num_requests=1,
            max_num_tokens=8192,
            kv_cache_manager=kv_cache_manager,
            request_ids=request_ids,
            prompt_lens=prompt_lens,
        )
        position_ids = [torch.arange(0, input_ids.size(-1), dtype=torch.int32)]
        position_ids = torch.cat(position_ids).unsqueeze(0).cuda()

        with torch.inference_mode():
            attn_metadata.prepare()
            logits = gemma3.forward(input_ids=input_ids,
                                    position_ids=position_ids,
                                    attn_metadata=attn_metadata)
            ref = hf_gemma3.forward(input_ids=input_ids.unsqueeze(0),
                                    position_ids=position_ids,
                                    past_key_values=hf_cache,
                                    use_cache=True)
            torch.testing.assert_close(logits,
                                       ref.logits[:, -1].float(),
                                       atol=0.05,
                                       rtol=0.05)

        # Generation phase.
        gen_input_ids = torch.tensor([900], dtype=torch.int, device=device)
        num_cached_tokens_per_seq = [input_ids.size(-1)]
        attn_metadata = metadata_cls(
            seq_lens=torch.tensor([gen_input_ids.size(-1)], dtype=torch.int),
            num_contexts=0,
            kv_cache_params=KVCacheParams(
                use_cache=True,
                num_cached_tokens_per_seq=num_cached_tokens_per_seq,
            ),
            kv_cache_manager=kv_cache_manager,
            request_ids=request_ids,
            prompt_lens=prompt_lens,
            max_num_requests=1,
            max_num_tokens=8192,
        )

        gen_position_ids = [
            torch.arange(input_ids.size(-1),
                         input_ids.size(-1) + gen_input_ids.size(-1))
        ]
        gen_position_ids = torch.cat(gen_position_ids).unsqueeze(0).cuda()
        with torch.inference_mode():
            attn_metadata.prepare()
            logits = gemma3.forward(input_ids=gen_input_ids,
                                    position_ids=gen_position_ids,
                                    attn_metadata=attn_metadata)
            ref = hf_gemma3.forward(input_ids=gen_input_ids.unsqueeze(0),
                                    position_ids=gen_position_ids,
                                    past_key_values=hf_cache,
                                    use_cache=True,
                                    cache_position=torch.IntTensor(
                                        [input_ids.size(-1)]).to(device),
                                    last_cache_position=input_ids.size(-1) + 1)
            torch.testing.assert_close(logits,
                                       ref.logits[:, -1].float(),
                                       atol=0.05,
                                       rtol=0.05)

        kv_cache_manager.shutdown()
