import json
import os
from pathlib import Path

from tests.ci.ci_register import register_cuda_ci

from miles.utils.external_utils import command_utils

register_cuda_ci(
    est_time=1000, suite="stage-c-8-gpu-b200", labels=["megatron", "model-scripts"], hardware=["blackwell"], num_gpus=8
)

MODEL_ORG = "Pinaster"
MODEL_NAME = "GLM-5.2_5layer"
MODEL_TYPE = "glm5.2-744B-A40B_5layer"
MEGATRON_MODEL_NAME = f"{MODEL_NAME}-megatron-dsa"
NVFP4_MODEL_NAME = f"{MODEL_NAME}-NVFP4-w4a16"
DSA_ARGS = "--megatron-to-hf-mode raw --dsa-impl megatron --dsa-kernel-backend cudnn "
NUM_GPUS = 8
ACTOR_NUM_GPUS = 4
ROLLOUT_NUM_GPUS = 4
ROLLOUT_GPUS_PER_ENGINE = 2
RUN_ID = command_utils.create_run_id()

MODEL_DIR = "/root/models"
DATA_DIR = "/root/datasets"
MEGATRON_PATH = "/root/TransformerEngine:/root/Megatron-LM"

EXTRA_HIGH_PRECISION_LAYERS_HF = (".shared_experts.",)
EXTRA_HIGH_PRECISION_LAYERS_MEGATRON = (
    ".shared_experts.linear_fc1",
    ".shared_experts.linear_fc2",
)

NVFP4_ENV = {
    "OPEN_TRAINING_NVFP4_FAKE_QAT_FLAG": "1",
    "SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16": "1",
    "SGLANG_FLASHINFER_MOE_FUSED_FINALIZE": "0",
    "NVTE_NVFP4_DISABLE_2D_QUANTIZATION": "1",
    "NVTE_NVFP4_DISABLE_RHT": "1",
    "NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING": "1",
    "NVTE_USE_FAST_MATH": "0",
    "NVTE_NVFP4_4OVER6": "all",
    "NVTE_NVFP4_4OVER6_E4M3_USE_256": "none",
    "NVTE_NVFP4_4OVER6_ERR_MODE": "MSE",
    "NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH": "1",
}

GLM5_ENV = {
    "SGLANG_DSA_FUSE_TOPK": "1",
    "SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD": "0",
    "SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK": "large",
    "INDEXER_ROPE_NEOX_STYLE": "0",
    "NVSHMEM_DISABLE_NCCL": "1",
}

# Fake QDQ supplies NVFP4 expert weights; TE runs BF16 forward and backward GEMMs.
TE_PRECISION_CONFIG = """
configs:
    bf16:
        transformer_engine_config_type: "TEQuantizationParams"
        training_recipe: {}
matchers:
    default_bf16:
        type: "glob"
        enabled: true
        pattern: "*"
        config: "bf16"
""".strip()


def _extra_high_precision_layers_hf_args() -> str:
    return "--extra-high-precision-layers-hf " + " ".join(EXTRA_HIGH_PRECISION_LAYERS_HF) + " "


def _extra_high_precision_layers_megatron_args() -> str:
    return "--extra-high-precision-layers-megatron " + " ".join(EXTRA_HIGH_PRECISION_LAYERS_MEGATRON) + " "


def _validate_glm_checkpoint():
    config_path = Path(MODEL_DIR) / MODEL_NAME / "config.json"
    if not config_path.exists():
        raise FileNotFoundError(f"{config_path} not found")

    with open(config_path) as f:
        config = json.load(f)

    if (
        config.get("model_type") != "glm_moe_dsa"
        or config.get("architectures") != ["GlmMoeDsaForCausalLM"]
        or config.get("num_hidden_layers") != 5
    ):
        raise RuntimeError(
            f"{config_path} must use native GLM-5.2 5-layer config with "
            f"model_type=glm_moe_dsa, architectures=[GlmMoeDsaForCausalLM], "
            "and num_hidden_layers=5"
        )
    if "auto_map" in config:
        raise RuntimeError(f"{config_path} must not contain auto_map. Try updating the checkpoint.")


def prepare():
    U = command_utils.default_config().create_backend()
    os.environ.update(NVFP4_ENV)
    U.exec_command_cpu(f"mkdir -p {MODEL_DIR} {DATA_DIR}")
    U.exec_command_cpu(f"hf download {MODEL_ORG}/{MODEL_NAME} --local-dir {MODEL_DIR}/{MODEL_NAME}")
    U.hf_download_dataset("zhuzilin/dapo-math-17k", data_dir=DATA_DIR)

    _validate_glm_checkpoint()
    U.exec_command_gpu_once(
        f"python tools/convert_hf_to_nvfp4.py "
        f"--model-dir {MODEL_DIR}/{MODEL_NAME} "
        f"--save-dir {MODEL_DIR}/{NVFP4_MODEL_NAME} "
        f"{_extra_high_precision_layers_hf_args()}",
        f"{MODEL_DIR}/{NVFP4_MODEL_NAME}",
    )

    U.convert_checkpoint(
        model_name=MEGATRON_MODEL_NAME,
        megatron_model_type=MODEL_TYPE,
        num_gpus_per_node=1,
        extra_args=(
            f"{DSA_ARGS}"
            "--tensor-model-parallel-size 1 "
            "--expert-tensor-parallel-size 1 "
            "--pipeline-model-parallel-size 1 "
            "--expert-model-parallel-size 1 "
        ),
        dir_dst=MODEL_DIR,
        hf_checkpoint=f"{MODEL_DIR}/{MODEL_NAME}",
        megatron_path=MEGATRON_PATH,
    )


def execute():
    U = command_utils.default_config().create_backend()
    os.environ.update(NVFP4_ENV)
    os.environ.update(GLM5_ENV)
    os.environ.setdefault("RAY_TMPDIR", "/tmp/ray")
    te_precision_config_path = command_utils.encode_pseudo_file(TE_PRECISION_CONFIG)

    ckpt_args = (
        f"--hf-checkpoint {MODEL_DIR}/{NVFP4_MODEL_NAME}/ " f"--ref-load {MODEL_DIR}/{MEGATRON_MODEL_NAME}_torch_dist "
    )

    rollout_args = (
        f"--prompt-data {DATA_DIR}/dapo-math-17k/dapo-math-17k.jsonl "
        "--input-key prompt "
        "--label-key label "
        "--apply-chat-template "
        "--rollout-shuffle "
        "--rm-type deepscaler "
        "--num-rollout 2 "
        "--rollout-batch-size 4 "
        "--n-samples-per-prompt 4 "
        "--rollout-max-response-len 100 "
        "--rollout-temperature 1 "
        "--global-batch-size 16 "
    )

    perf_args = (
        "--tensor-model-parallel-size 1 "
        # Let the STE propagate gradients to the original expert parameters.
        "--no-gradient-accumulation-fusion "
        "--pipeline-model-parallel-size 1 "
        f"--context-parallel-size {ACTOR_NUM_GPUS} "
        f"--expert-model-parallel-size {ACTOR_NUM_GPUS} "
        "--expert-tensor-parallel-size 1 "
        "--recompute-granularity full "
        "--recompute-method uniform "
        "--recompute-num-layers 1 "
        "--use-dynamic-batch-size "
        "--max-tokens-per-gpu 2048 "
        "--data-pad-size-multiplier 1024 "
        "--log-probs-chunk-size 16384 "
    )

    grpo_args = (
        "--advantage-estimator grpo "
        "--use-kl-loss "
        "--kl-loss-coef 0.00 "
        "--kl-loss-type low_var_kl "
        "--kl-coef 0.00 "
        "--entropy-coef 0.00 "
        "--eps-clip 0.2 "
        "--eps-clip-high 0.28 "
        "--use-tis "
        "--tis-clip-low 0.5 "
        "--tis-clip 2.0 "
    )

    optimizer_args = (
        "--optimizer adam "
        "--lr 1e-6 "
        "--lr-decay-style constant "
        "--weight-decay 0.1 "
        "--adam-beta1 0.9 "
        "--adam-beta2 0.98 "
        "--optimizer-cpu-offload "
        "--overlap-cpu-optimizer-d2h-h2d "
        "--use-precision-aware-optimizer "
    )

    sglang_args = (
        "--sglang-mem-fraction-static 0.7 "
        "--sglang-enable-dp-attention "
        "--sglang-attention-backend nsa "
        "--sglang-dsa-decode-backend flashmla_kv "
        "--sglang-dsa-prefill-backend flashmla_sparse "
        "--sglang-dsa-topk-backend flashinfer "
        "--sglang-kv-cache-dtype fp8_e4m3 "
        "--sglang-page-size 64 "
        f"--rollout-num-gpus-per-engine {ROLLOUT_GPUS_PER_ENGINE} "
        "--sglang-moe-runner-backend flashinfer_cutedsl "
        "--sglang-moe-a2a-backend none "
        f"--sglang-ep-size {ROLLOUT_GPUS_PER_ENGINE} "
        f"--sglang-dp-size {ROLLOUT_GPUS_PER_ENGINE} "
        "--sglang-moe-dense-tp-size 1 "
        "--sglang-enable-dp-lm-head "
        "--sglang-enable-fp32-lm-head "
        "--sglang-cuda-graph-max-bs-decode 256 "
        "--sglang-max-running-requests 512 "
        f"--sglang-chunked-prefill-size {2048 * ROLLOUT_GPUS_PER_ENGINE} "
        "--sglang-watchdog-timeout 3600 "
    )

    ci_args = "--ci-test --ci-disable-logprobs-checker --ci-disable-weight-update-checker --ci-disable-kl-checker "

    mixed_precision_args = (
        "--transformer-impl transformer_engine "
        "--bf16 "
        f"{_extra_high_precision_layers_hf_args()}"
        f"{_extra_high_precision_layers_megatron_args()}"
        f"--te-precision-config-file {te_precision_config_path} "
    )

    misc_args = (
        f"{DSA_ARGS}"
        "--qkv-format thd "
        "--use-rollout-routing-replay "
        "--use-miles-router "
        "--sglang-disable-shared-experts-fusion "
        "--attention-dropout 0.0 "
        "--hidden-dropout 0.0 "
        "--accumulate-allreduce-grads-in-fp32 "
        "--attention-softmax-in-fp32 "
        "--moe-router-use-torch-mm "
        "--attention-backend flash "
        "--cp-comm-type allgather "
        "--miles-dsa-topk-backend flashinfer "
        "--update-weight-transfer-mode broadcast_packed "
        f"--update-weight-buffer-size {2 * 1024 ** 3} "
        "--actor-num-nodes 1 "
        f"--actor-num-gpus-per-node {ACTOR_NUM_GPUS} "
        f"--num-gpus-per-node {NUM_GPUS} "
        f"--rollout-num-gpus {ROLLOUT_NUM_GPUS} "
        "--moe-enable-deepep "
        "--moe-token-dispatcher-type flex "
        # Event logging requests weight checksums that SGLang does not support for NVFP4.
        f"--save-debug-rollout-data /root/shared_data/{RUN_ID}/dump_details/rollout_data/{{rollout_id}}.pt "
        f"--save-debug-train-data /root/shared_data/{RUN_ID}/dump_details/train_data/{{rollout_id}}_{{rank}}.pt "
        f"--save-debug-trajectory-data /root/shared_data/{RUN_ID}/dump_details/trajectory/{{rollout_id}}.jsonl "
    )

    train_args = (
        f"{ckpt_args} "
        f"{rollout_args} "
        f"{optimizer_args} "
        f"{grpo_args} "
        f"{command_utils.get_default_wandb_args(__file__, run_id=RUN_ID)} "
        f"{perf_args} "
        f"{sglang_args} "
        f"{ci_args} "
        f"{mixed_precision_args} "
        f"{misc_args} "
    )

    U.execute_train(
        train_args=train_args,
        num_gpus_per_node=NUM_GPUS,
        megatron_model_type=MODEL_TYPE,
        megatron_path=MEGATRON_PATH,
        extra_env_vars={**NVFP4_ENV, **GLM5_ENV},
    )


if __name__ == "__main__":
    prepare()
    for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
        os.environ.pop(proxy_var, None)
    execute()