"""vLLM serving tests for Nemotron-3.

Nemotron-3 uses 'backbone.*' weight prefix in HF checkpoints, but vLLM
remaps to 'model.*' internally via WeightsMapper. The adapter conversion
applies the same remap so PEFT keys match vLLM's parameter names.

Nemotron-3 is a hybrid Mamba+Attention MoE architecture.

Run individual variants:
    CUDA_VISIBLE_DEVICES=0,1 .../python -m pytest .../test_nemotron.py::TestNemotron3Nano -v -s
    CUDA_VISIBLE_DEVICES=0,1,2,3 .../python -m pytest .../test_nemotron.py::TestNemotron3Super -v -s
"""

from __future__ import annotations

import pytest
import torch
from vllm import LLM
from vllm.lora.request import LoRARequest

from tests.weights.vllm_serving.conftest import (
    LORA_RANK,
    PROMPT,
    convert_and_load,
    generate,
    load_hf_config_dict,
    save_tinker_adapter,
)


def _make_nemotron_adapter(model: str) -> tuple[dict, dict[str, torch.Tensor]]:
    """Create a synthetic Tinker adapter for a Nemotron attention layer only.

    Returns (config_dict, adapter_weights).
    """
    config = load_hf_config_dict(model)
    hidden = config["hidden_size"]
    num_heads = config["num_attention_heads"]
    head_dim = config.get("head_dim", hidden // num_heads)
    q_dim = num_heads * head_dim

    # Find an attention layer. Nemotron-3 uses hybrid_override_pattern:
    # M=Mamba, E=MoE, *=Attention
    pattern = config.get("hybrid_override_pattern", "")
    attn_idx = next((i for i, ch in enumerate(pattern) if ch == "*"), None)
    if attn_idx is None:
        pytest.skip("No attention layer found in Nemotron config")

    # Tinker adapter uses backbone.* prefix (matching HF checkpoint)
    prefix = f"base_model.model.backbone.layers.{attn_idx}.mixer"
    weights = {
        f"{prefix}.q_proj.lora_A.weight": torch.randn(LORA_RANK, hidden) * 0.01,
        f"{prefix}.q_proj.lora_B.weight": torch.randn(q_dim, LORA_RANK) * 0.01,
    }
    return config, weights


def _make_nemotron_moe_adapter(model: str) -> tuple[dict, dict[str, torch.Tensor]]:
    """Create a synthetic Tinker adapter with attention + Mamba + MoE expert LoRA.

    Extends :func:`_make_nemotron_adapter` with:
    - Mamba layer: gate_proj/x_proj LoRA (merged into in_proj during conversion)
    - MoE layer: w1 (up_proj) + w2 (down_proj) 3D LoRA with broadcast, plus empty w3

    Returns (config_dict, adapter_weights).
    """
    config, weights = _make_nemotron_adapter(model)
    hidden = config["hidden_size"]
    n_experts = config.get("n_routed_experts", 128)
    moe_intermediate = config.get("moe_intermediate_size", 1856)
    mamba_intermediate = config.get("mamba_num_heads", 64) * config.get("mamba_head_dim", 64)

    pattern = config.get("hybrid_override_pattern", "")
    mamba_idx = next((i for i, ch in enumerate(pattern) if ch == "M"), None)
    moe_idx = next((i for i, ch in enumerate(pattern) if ch == "E"), None)
    if mamba_idx is None:
        pytest.skip("No Mamba layer found in Nemotron config")
    if moe_idx is None:
        pytest.skip("No MoE layer found in Nemotron config")

    mamba_prefix = f"base_model.model.backbone.layers.{mamba_idx}.mixer"
    moe_prefix = f"base_model.model.backbone.layers.{moe_idx}.mixer"
    weights.update(
        {
            # Mamba layer: gate_proj/x_proj trained separately by Tinker,
            # merged into fused in_proj during adapter conversion.
            f"{mamba_prefix}.gate_proj.lora_A.weight": torch.randn(LORA_RANK, hidden) * 0.01,
            f"{mamba_prefix}.gate_proj.lora_B.weight": torch.randn(mamba_intermediate, LORA_RANK)
            * 0.01,
            f"{mamba_prefix}.x_proj.lora_A.weight": torch.randn(LORA_RANK, hidden) * 0.01,
            f"{mamba_prefix}.x_proj.lora_B.weight": torch.randn(mamba_intermediate, LORA_RANK)
            * 0.01,
            # Expert w1 (up_proj): lora_A shared (1), lora_B per-expert
            f"{moe_prefix}.experts.w1.lora_A.weight": torch.randn(1, LORA_RANK, hidden) * 0.01,
            f"{moe_prefix}.experts.w1.lora_B.weight": torch.randn(
                n_experts, moe_intermediate, LORA_RANK
            )
            * 0.01,
            # Expert w2 (down_proj): lora_A per-expert, lora_B shared (1)
            f"{moe_prefix}.experts.w2.lora_A.weight": torch.randn(
                n_experts, LORA_RANK, moe_intermediate
            )
            * 0.01,
            f"{moe_prefix}.experts.w2.lora_B.weight": torch.randn(1, hidden, LORA_RANK) * 0.01,
            # Expert w3 (gate_proj): empty — Nemotron has no gate_proj
            f"{moe_prefix}.experts.w3.lora_A.weight": torch.empty(0),
            f"{moe_prefix}.experts.w3.lora_B.weight": torch.empty(0),
        }
    )
    return config, weights


def _verify_nemotron_peft_keys(peft_weights: dict[str, torch.Tensor], peft_config: dict) -> None:
    """Assert PEFT keys use model.* (not backbone.*) and target_modules is correct."""
    assert not any("backbone" in k for k in peft_weights), (
        "PEFT keys should not contain 'backbone' — vLLM remaps to 'model'"
    )
    assert any("model.layers" in k for k in peft_weights), (
        "PEFT keys should use 'model.layers' prefix"
    )
    assert "q_proj" in peft_config["target_modules"]


class TestNemotron3Nano:
    """Nemotron-3-Nano-30B-A3B: 30B MoE (3B active), TP=2."""

    MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"

    @pytest.fixture(scope="class")
    def llm(self):
        return LLM(
            model=self.MODEL,
            enable_lora=True,
            max_lora_rank=LORA_RANK * 2,  # doubled for fused in_proj merge
            max_loras=1,
            max_model_len=256,
            enforce_eager=True,
            gpu_memory_utilization=0.5,
            tensor_parallel_size=2,
            trust_remote_code=True,
        )

    def test_attention_adapter(self, llm, tmp_path):
        """Verify backbone→model remap and vLLM serving for Nano variant."""
        _, adapter_weights = _make_nemotron_adapter(self.MODEL)

        adapter_path = tmp_path / "tinker_adapter"
        peft_path = tmp_path / "peft_adapter"
        save_tinker_adapter(adapter_path, adapter_weights)
        peft_weights, peft_config = convert_and_load(self.MODEL, adapter_path, peft_path)

        _verify_nemotron_peft_keys(peft_weights, peft_config)

        base_text = generate(llm, PROMPT)
        assert len(base_text) > 0

        lora_req = LoRARequest("nemotron3_nano", 1, str(peft_path))
        lora_text = generate(llm, PROMPT, lora_request=lora_req)
        assert len(lora_text) > 0

        print(f"\n  PEFT keys: {sorted(peft_weights.keys())}")
        print(f"  Base: {base_text!r}")
        print(f"  LoRA: {lora_text!r}")

    def test_moe_expert_adapter(self, llm, tmp_path):
        """Verify Mamba + MoE expert LoRA converts and serves correctly.

        Regression test for #547: build_lora_adapter crashed on empty expert
        LoRA tensors, mapped w1→gate_proj instead of w1→up_proj, and didn't
        merge Mamba gate_proj/x_proj into fused in_proj.
        """
        _, adapter_weights = _make_nemotron_moe_adapter(self.MODEL)

        adapter_path = tmp_path / "tinker_moe_adapter"
        peft_path = tmp_path / "peft_moe_adapter"
        save_tinker_adapter(adapter_path, adapter_weights)
        peft_weights, peft_config = convert_and_load(self.MODEL, adapter_path, peft_path)

        # Verify backbone→model prefix remap
        assert not any("backbone" in k for k in peft_weights), (
            "PEFT keys should not contain 'backbone'"
        )

        # Verify expert keys use up_proj/down_proj (not gate_proj)
        expert_keys = [k for k in peft_weights if "experts" in k]
        assert len(expert_keys) > 0, "Should have per-expert PEFT keys"
        assert not any("gate_proj" in k for k in expert_keys), (
            "Expert keys should use up_proj/down_proj, not gate_proj"
        )
        assert any("up_proj" in k for k in expert_keys)
        assert any("down_proj" in k for k in expert_keys)

        # Verify Mamba gate_proj/x_proj merged into in_proj
        assert not any("gate_proj" in k for k in peft_weights if "experts" not in k), (
            "Mamba gate_proj should be merged into in_proj"
        )
        assert not any("x_proj" in k for k in peft_weights), (
            "Mamba x_proj should be merged into in_proj"
        )
        assert any("in_proj" in k for k in peft_weights), "Merged in_proj LoRA should be present"
        assert "in_proj" in peft_config["target_modules"]
        assert peft_config["rank_pattern"].get("in_proj", 0) > peft_config["r"], (
            "in_proj should have doubled rank in rank_pattern"
        )

        # Verify vLLM can load and generate with the full adapter
        base_text = generate(llm, PROMPT)
        assert len(base_text) > 0

        lora_req = LoRARequest("nemotron3_nano_moe", 1, str(peft_path))
        lora_text = generate(llm, PROMPT, lora_request=lora_req)
        assert len(lora_text) > 0

        print(f"\n  Expert PEFT keys: {len(expert_keys)}")
        print(f"  target_modules: {peft_config['target_modules']}")
        print(f"  rank_pattern: {peft_config['rank_pattern']}")
        print(f"  Base: {base_text!r}")
        print(f"  LoRA: {lora_text!r}")


class TestNemotron3Super:
    """Nemotron-3-Super-120B-A12B: 120B MoE (12B active), TP=4."""

    MODEL = "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16"

    @pytest.fixture(scope="class")
    def llm(self):
        return LLM(
            model=self.MODEL,
            enable_lora=True,
            max_lora_rank=LORA_RANK * 2,  # doubled for fused in_proj merge
            max_loras=1,
            max_model_len=256,
            enforce_eager=True,
            gpu_memory_utilization=0.5,
            tensor_parallel_size=4,
            trust_remote_code=True,
        )

    def test_attention_adapter(self, llm, tmp_path):
        """Verify backbone→model remap and vLLM serving for Super variant."""
        _, adapter_weights = _make_nemotron_adapter(self.MODEL)

        adapter_path = tmp_path / "tinker_adapter"
        peft_path = tmp_path / "peft_adapter"
        save_tinker_adapter(adapter_path, adapter_weights)
        peft_weights, peft_config = convert_and_load(self.MODEL, adapter_path, peft_path)

        _verify_nemotron_peft_keys(peft_weights, peft_config)

        base_text = generate(llm, PROMPT)
        assert len(base_text) > 0

        lora_req = LoRARequest("nemotron3_super", 1, str(peft_path))
        lora_text = generate(llm, PROMPT, lora_request=lora_req)
        assert len(lora_text) > 0

        print(f"\n  PEFT keys: {sorted(peft_weights.keys())}")
        print(f"  Base: {base_text!r}")
        print(f"  LoRA: {lora_text!r}")

    def test_moe_expert_adapter(self, llm, tmp_path):
        """Verify Mamba + MoE expert LoRA converts and serves correctly.

        Regression test for #547 on the Super variant.
        """
        _, adapter_weights = _make_nemotron_moe_adapter(self.MODEL)

        adapter_path = tmp_path / "tinker_moe_adapter"
        peft_path = tmp_path / "peft_moe_adapter"
        save_tinker_adapter(adapter_path, adapter_weights)
        peft_weights, peft_config = convert_and_load(self.MODEL, adapter_path, peft_path)

        # Verify backbone→model prefix remap
        assert not any("backbone" in k for k in peft_weights), (
            "PEFT keys should not contain 'backbone'"
        )

        # Verify expert keys use up_proj/down_proj (not gate_proj)
        expert_keys = [k for k in peft_weights if "experts" in k]
        assert len(expert_keys) > 0, "Should have per-expert PEFT keys"
        assert not any("gate_proj" in k for k in expert_keys), (
            "Expert keys should use up_proj/down_proj, not gate_proj"
        )
        assert any("up_proj" in k for k in expert_keys)
        assert any("down_proj" in k for k in expert_keys)

        # Verify Mamba gate_proj/x_proj merged into in_proj
        assert any("in_proj" in k for k in peft_weights), "Merged in_proj LoRA should be present"

        # Verify vLLM can load and generate with the full adapter
        base_text = generate(llm, PROMPT)
        assert len(base_text) > 0

        lora_req = LoRARequest("nemotron3_super_moe", 1, str(peft_path))
        lora_text = generate(llm, PROMPT, lora_request=lora_req)
        assert len(lora_text) > 0

        print(f"\n  Expert PEFT keys: {len(expert_keys)}")
        print(f"  target_modules: {peft_config['target_modules']}")
        print(f"  rank_pattern: {peft_config['rank_pattern']}")
        print(f"  Base: {base_text!r}")
        print(f"  LoRA: {lora_text!r}")
