"""
Lightweight inference runtime (replaces Megatron/DeepSpeed training stack in release builds).
"""
from __future__ import annotations

import math
import os
from typing import Any, Dict, Optional

import torch
import torch.distributed as dist

_GLOBAL_ARGS: Optional[Any] = None


def set_global_args(args: Any) -> None:
    global _GLOBAL_ARGS
    _GLOBAL_ARGS = args


def get_args() -> Any:
    if _GLOBAL_ARGS is None:
        raise RuntimeError("Inference args not initialized. Call set_global_args() first.")
    return _GLOBAL_ARGS


def print_rank_0(msg: str) -> None:
    if not dist.is_initialized() or dist.get_rank() == 0:
        print(msg, flush=True)


class _NoOpTimer:
    def start(self, *args, **kwargs):
        return self

    def stop(self, *args, **kwargs):
        return self

    def __enter__(self):
        return self

    def __exit__(self, *args):
        return False


def get_timers():
    return _NoOpTimer()


def get_model_parallel_world_size() -> int:
    return 1


def get_tensor_model_parallel_rank() -> int:
    return 0


def get_data_parallel_world_size() -> int:
    # When sequence parallelism is on, the data-parallel axis is the complement
    # of the SP axis (each SP group works on ONE sample together).
    from .ulysses import sequence_parallel_is_enabled, get_dp_world_size
    if sequence_parallel_is_enabled():
        return get_dp_world_size()
    if dist.is_initialized():
        return dist.get_world_size()
    return 1


def get_data_parallel_rank() -> int:
    from .ulysses import sequence_parallel_is_enabled, get_dp_rank
    if sequence_parallel_is_enabled():
        return get_dp_rank()
    if dist.is_initialized():
        return dist.get_rank()
    return 0


def broadcast_data(keys, data: Dict[str, Any], dtype: torch.dtype):
    if get_model_parallel_world_size() <= 1:
        return data
    return data


class MPU:
    get_model_parallel_world_size = staticmethod(get_model_parallel_world_size)
    get_tensor_model_parallel_rank = staticmethod(get_tensor_model_parallel_rank)
    get_data_parallel_world_size = staticmethod(get_data_parallel_world_size)
    get_data_parallel_rank = staticmethod(get_data_parallel_rank)
    broadcast_data = staticmethod(broadcast_data)


mpu = MPU()


def init_method_normal(std: float):
    def init_(tensor):
        return torch.nn.init.normal_(tensor, mean=0.0, std=std)

    return init_


def scaled_init_method_normal(std: float, num_layers: int):
    scale = std / math.sqrt(2.0 * num_layers)

    def init_(tensor):
        return torch.nn.init.normal_(tensor, mean=0.0, std=scale)

    return init_


def init_distributed(backend: str = "nccl") -> tuple[int, int, int]:
    if dist.is_initialized():
        rank = dist.get_rank()
        world_size = dist.get_world_size()
    else:
        # When launched with plain `python` (single process, no torchrun), the
        # standard env:// variables are absent. Fill in single-process defaults
        # so init_process_group does not hang waiting for a rendezvous.
        os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
        os.environ.setdefault("MASTER_PORT", "29500")
        os.environ.setdefault("RANK", "0")
        os.environ.setdefault("WORLD_SIZE", "1")
        os.environ.setdefault("LOCAL_RANK", "0")
        dist.init_process_group(backend=backend, init_method="env://")
        rank = dist.get_rank()
        world_size = dist.get_world_size()
    local_rank = int(os.environ.get("LOCAL_RANK", rank % max(torch.cuda.device_count(), 1)))
    torch.cuda.set_device(local_rank)
    return rank, local_rank, world_size