"""Shared torch runtime helpers for device/DDP/DataParallel setup."""

from __future__ import annotations

import inspect
from contextlib import nullcontext
from typing import Callable, Optional, Tuple

import torch
import torch.nn as nn
from torch.amp import GradScaler, autocast
from torch.nn.parallel import DistributedDataParallel as DDP

from ins_pricing.modelling.bayesopt.utils.distributed_utils import DistributedUtils
from ins_pricing.utils.device import DeviceManager


def setup_ddp_if_requested(use_ddp: bool) -> Tuple[bool, int, int, int]:
    """Initialize DDP only when requested."""
    if not bool(use_ddp):
        return False, 0, 0, 1
    ddp_ok, local_rank, rank, world_size = DistributedUtils.setup_ddp()
    return bool(ddp_ok), int(local_rank), int(rank), int(world_size)


def resolve_training_device(
    *,
    is_ddp_enabled: bool,
    local_rank: int,
    use_gpu: bool = True,
) -> torch.device:
    """Resolve training device with stable fallback order: cuda > mps > cpu."""
    device = DeviceManager.resolve_training_device(
        is_ddp_enabled=is_ddp_enabled,
        local_rank=local_rank,
        use_gpu=use_gpu,
    )
    if device is None:
        return torch.device("cpu")
    return device


def _supports_autocast(device_type: str) -> bool:
    checker = getattr(torch.amp.autocast_mode, "is_autocast_available", None)
    if callable(checker):
        try:
            return bool(checker(device_type))
        except Exception:
            return False
    return device_type == "cuda"


def _grad_scaler_accepts_device_kw() -> bool:
    try:
        return "device" in inspect.signature(GradScaler).parameters
    except Exception:
        return False


def create_autocast_context(device_type: str, enabled: Optional[bool] = None):
    """Create autocast context with safe fallback when backend is unsupported."""
    if enabled is None:
        enabled = device_type in ("cuda", "mps")
    enabled = bool(enabled) and _supports_autocast(device_type)
    if not enabled:
        return nullcontext()
    if device_type == "mps":
        return autocast(device_type=device_type, dtype=torch.float16, enabled=True)
    return autocast(device_type=device_type, enabled=True)


def create_grad_scaler(device_type: str, enabled: Optional[bool] = None) -> GradScaler:
    """Create AMP GradScaler with backend-aware compatibility fallback."""
    if enabled is None:
        enabled = device_type in ("cuda", "mps")
    enabled = bool(enabled)
    target_device = device_type if device_type in {"cuda", "mps", "cpu", "xpu"} else "cuda"

    if _grad_scaler_accepts_device_kw():
        return GradScaler(device=target_device, enabled=enabled)

    # Legacy API fallback: only CUDA is safe without a device argument.
    if target_device != "cuda":
        return GradScaler(enabled=False)
    return GradScaler(enabled=enabled)


def wrap_model_for_parallel(
    core: nn.Module,
    *,
    device: torch.device,
    use_data_parallel: bool,
    use_ddp_requested: bool,
    is_ddp_enabled: bool,
    local_rank: int,
    ddp_find_unused_parameters: bool = False,
    fallback_log: Optional[Callable[[str], None]] = None,
) -> Tuple[nn.Module, bool, torch.device]:
    """Wrap model with DDP or DataParallel based on resolved runtime flags."""
    if is_ddp_enabled:
        core = core.to(device)
        core = DDP(
            core,
            device_ids=[int(local_rank)],
            output_device=int(local_rank),
            find_unused_parameters=bool(ddp_find_unused_parameters),
        )
        return core, False, device

    should_use_dp = (
        bool(use_data_parallel)
        and device.type == "cuda"
        and torch.cuda.device_count() > 1
    )
    if should_use_dp:
        if bool(use_ddp_requested) and fallback_log is not None:
            fallback_log(
                ">>> DDP requested but not initialized; falling back to DataParallel."
            )
        core = nn.DataParallel(core, device_ids=list(range(torch.cuda.device_count())))
        return core, True, torch.device("cuda")

    return core, False, device


__all__ = [
    "setup_ddp_if_requested",
    "resolve_training_device",
    "create_autocast_context",
    "create_grad_scaler",
    "wrap_model_for_parallel",
]
