import logging

import pytest
import torch
import torch.distributed as dist
from packaging import version
from transformers import AutoModelForCausalLM

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

SEED = 0xdeadbeef


def pytest_addoption(parser):
    parser.addoption(
        "--measure-perf",
        action="store_true",
        default=False,
        help=
        "Measure execution time and peak memory usage during optimizer step.",
    )

    parser.addoption(
        "--do-profile",
        action="store_true",
        default=False,
        help="Enable profiling during tests.",
    )

    parser.addoption(
        "--skip-verify",
        action="store_true",
        default=False,
        help=
        "Skip verification of optimizer step correctness with sequential implementation.\n"
        "This can be useful when GPU memory is limited.",
    )


def pytest_configure(config):
    if config.getoption(
            "--do-profile") and not config.getoption("--measure-perf"):
        raise pytest.UsageError(
            "--do-profile requires --measure-perf. Please enable both flags.")


@pytest.fixture(scope="session")
def measure_perf(request):
    return request.config.getoption("--measure-perf")


@pytest.fixture(scope="session")
def do_profile(request):
    return request.config.getoption("--do-profile")


@pytest.fixture(scope="session")
def skip_verify(request):
    return request.config.getoption("--skip-verify")


@pytest.fixture(scope="session", autouse=True)
def init_dist(request):
    if version.parse(torch.__version__) < version.parse("2.8"):
        pytest.skip("torch>=2.8.0 is required for parallel muon")
        return

    try:
        dist.init_process_group(backend="nccl")
        torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count())
    except Exception as e:
        print(f"Failed to initialize torch.distributed: {e}")
        pytest.skip("Failed to initialize torch.distributed")

    if dist.get_world_size() != 8:
        pytest.skip("Need 8 processes in dist group. "
                    "You can run with `torchrun --nproc-per-node=8 "
                    "--local-ranks-filter 0 -m pytest "
                    "test_rms_norm_sequence_parallel.py`."
                    "To run with less than 8 gpus, modify "
                    "the test cases accordingly.")

    yield
    dist.destroy_process_group()


@pytest.fixture(scope="session")
def inputs():
    """Load Motif-2.6B model and generate random gradients for testing.
    Returns:
        tuple[torch.nn.Module, list[torch.Tensor], dict[int, torch.Tensor]]:
        - torch.nn.Module: The Motif-2.6B model.
        - list[torch.Tensor]: A list of random gradients for each model parameter.
        - dict[int, torch.Tensor]: A dictionary mapping layer indices to random QK logits.
    """
    model_name = "Motif-Technologies/Motif-2.6B-4layer-random"

    torch.manual_seed(SEED)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(SEED)

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
    )
    logger.info(
        f"Loaded model {model_name}. ({len(list(model.parameters()))} parameters)"
    )

    grads: list[torch.Tensor] = []
    for param in model.parameters():
        grad = torch.randn_like(param, device=param.device, dtype=param.dtype)
        grads.append(grad)

    qk_logits: dict[int, torch.Tensor] = {
        i:
        torch.randn(model.config.num_attention_heads,
                    device=model.device,
                    dtype=torch.bfloat16)
        for i in range(model.config.num_hidden_layers)
    }

    return [model, grads, qk_logits]