"""Distributed-training config builders — Accelerate FSDP + DeepSpeed ZeRO.

Pure-Python: returns plain dicts that the caller writes to YAML/JSON for
`accelerate launch --config_file <path>` or
`deepspeed --deepspeed_config <path>`.

Hard invariant: MI300X xGMI permits only 1- or 8-GPU FSDP. The 2/4-GPU
configurations have a known bandwidth bug; this module rejects them.
"""

from __future__ import annotations

from typing import Literal


def build_fsdp_config(
    num_gpus: Literal[1, 8],
    *,
    shard_size: Literal["FULL_SHARD", "SHARD_GRAD_OP", "NO_SHARD"] = "FULL_SHARD",
    transformer_layer_class: str = "Qwen3DecoderLayer",
    mixed_precision: Literal["no", "fp16", "bf16"] = "bf16",
) -> dict[str, object]:
    """Return an Accelerate FSDP config dict for `num_gpus`."""
    if num_gpus not in (1, 8):
        msg = f"MI300X xGMI permits only 1 or 8 GPUs; got {num_gpus}."
        raise ValueError(msg)
    return {
        "compute_environment": "LOCAL_MACHINE",
        "distributed_type": "FSDP" if num_gpus > 1 else "NO",
        "downcast_bf16": "no",
        "machine_rank": 0,
        "main_training_function": "main",
        "mixed_precision": mixed_precision,
        "num_machines": 1,
        "num_processes": num_gpus,
        "rdzv_backend": "static",
        "same_network": True,
        "tpu_env": [],
        "tpu_use_cluster": False,
        "tpu_use_sudo": False,
        "use_cpu": False,
        "fsdp_config": {
            "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
            "fsdp_backward_prefetch_policy": "BACKWARD_PRE",
            "fsdp_forward_prefetch": False,
            "fsdp_offload_params": False,
            "fsdp_sharding_strategy": shard_size,
            "fsdp_state_dict_type": "FULL_STATE_DICT",
            "fsdp_sync_module_states": True,
            "fsdp_transformer_layer_cls_to_wrap": transformer_layer_class,
            "fsdp_use_orig_params": True,
        },
    }


def build_deepspeed_config(
    *,
    zero_stage: Literal[1, 2, 3] = 3,
    offload_optimizer: bool = False,
    offload_param: bool = False,
    overlap_comm: bool = True,
) -> dict[str, object]:
    """Return a DeepSpeed ZeRO config dict at the given stage."""
    if zero_stage not in (1, 2, 3):
        msg = f"zero_stage must be 1/2/3; got {zero_stage}"
        raise ValueError(msg)
    return {
        "bf16": {"enabled": True},
        "zero_optimization": {
            "stage": zero_stage,
            "offload_optimizer": {"device": "cpu" if offload_optimizer else "none"},
            "offload_param": {"device": "cpu" if offload_param else "none"},
            "overlap_comm": overlap_comm,
            "contiguous_gradients": True,
            "reduce_bucket_size": "auto",
            "stage3_prefetch_bucket_size": "auto",
            "stage3_param_persistence_threshold": "auto",
        },
        "gradient_accumulation_steps": "auto",
        "gradient_clipping": "auto",
        "steps_per_print": 100,
        "train_batch_size": "auto",
        "train_micro_batch_size_per_gpu": "auto",
        "wall_clock_breakdown": False,
    }


__all__ = ["build_deepspeed_config", "build_fsdp_config"]