#!/usr/bin/env python3
"""
Daimon Training Template — Validation Tests
=============================================

Run these BEFORE launching training to catch configuration issues early.
Each test is independent and reports PASS/FAIL.

Usage:
    python3 /workspace/runpod-template/test_template.py

These tests run on a SINGLE GPU (no distributed launch needed).
"""

import os
import sys
import json
import time
import traceback

RESULTS = []


def test(name):
    """Decorator to register and run a test."""
    def decorator(fn):
        def wrapper():
            try:
                fn()
                RESULTS.append(("PASS", name, None))
                print(f"  PASS: {name}")
            except Exception as e:
                RESULTS.append(("FAIL", name, str(e)))
                print(f"  FAIL: {name}")
                print(f"        {e}")
                traceback.print_exc()
        wrapper.__name__ = name
        wrapper._test = True
        return wrapper
    return decorator


# ── Test 1: GPU with sufficient VRAM ──────────────────────────────────────

@test("GPU with >= 140GB VRAM is available")
def test_gpu_vram():
    import torch
    gpu_count = torch.cuda.device_count()
    assert gpu_count >= 1, (
        f"No GPUs detected. Need at least 1x H200 SXM 141GB."
    )

    # Find the GPU with the most VRAM
    max_vram_gb = 0
    for i in range(gpu_count):
        name = torch.cuda.get_device_name(i)
        mem_gb = torch.cuda.get_device_properties(i).total_memory / 1e9
        max_vram_gb = max(max_vram_gb, mem_gb)
        print(f"        GPU {i}: {name}, {mem_gb:.1f} GB")

    assert max_vram_gb >= 140, (
        f"Largest GPU has {max_vram_gb:.1f} GB VRAM. Need >= 140 GB (H200 SXM). "
        f"Full SFT needs ~90GB GPU (70GB model + 20GB activations)."
    )


# ── Test 2: System RAM >= 180GB (critical for CPU offload) ───────────────

@test("System RAM >= 180GB for CPU-offloaded optimizer")
def test_system_ram():
    """
    Full-parameter SFT offloads gradients (~70GB) and Adafactor states (~35GB)
    to CPU RAM. Without enough system RAM, training will OOM on the CPU side.
    """
    ram_bytes = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")
    ram_gb = ram_bytes / 1e9
    print(f"        System RAM: {ram_gb:.0f} GB")

    assert ram_gb >= 180, (
        f"System RAM is {ram_gb:.0f} GB. Need >= 180 GB. "
        f"CPU-offloaded memory budget: gradients (~70GB) + Adafactor (~35GB) = ~105GB, "
        f"plus OS and data loading overhead. "
        f"AdamW would need ~280GB — that's why we use Adafactor."
    )

    # Warn if tight
    if ram_gb < 200:
        print(f"        WARNING: {ram_gb:.0f}GB is tight. 200GB+ recommended.")
        print(f"        CPU budget: ~105GB for offloaded states + ~30GB overhead")


# ── Test 3: Full SFT config is valid (no LoRA) ──────────────────────────

@test("Full SFT config is valid (no LoRA)")
def test_sft_config():
    import yaml

    config_paths = [
        "/workspace/runpod-template/train_daimon_config.yaml",
        os.path.join(os.path.dirname(__file__), "train_daimon_config.yaml"),
    ]
    found = None
    for p in config_paths:
        if os.path.exists(p):
            found = p
            break

    assert found is not None, (
        f"Config not found. Looked in: {config_paths}"
    )

    with open(found) as f:
        config = yaml.safe_load(f)

    # Verify NO LoRA section — this is full-parameter SFT
    assert "lora" not in config, (
        "Config still has a 'lora' section. This template does full-parameter SFT — "
        "remove the lora section entirely."
    )

    # Verify Adafactor optimizer is configured
    optimizer = config.get("optimizer", "")
    assert optimizer.lower() == "adafactor", (
        f"Optimizer must be 'adafactor' for full SFT on single node. "
        f"Got: '{optimizer}'. AdamW needs ~280GB CPU RAM — not viable."
    )

    # Verify DeepSpeed config is referenced
    ds_config = config.get("deepspeed_config")
    assert ds_config is not None, (
        "Config must reference deepspeed_config for ZeRO-2 CPU offload. "
        "Full-parameter SFT cannot fit without gradient sharding."
    )

    # Verify learning rate is appropriate for full SFT (not LoRA rate)
    lr = config.get("learning_rate", 0)
    assert lr <= 1e-4, (
        f"Learning rate {lr} is too high for full SFT. "
        f"LoRA uses 2e-4, but full SFT should be 1e-5 to 5e-6 to avoid "
        f"destabilizing MoE routing gates."
    )

    # Verify model revision is pinned
    revision = config.get("model_revision")
    assert revision is not None and len(revision) >= 10, (
        f"model_revision should be pinned to a specific commit hash. Got: {revision}"
    )

    # Verify bf16 is enabled
    assert config.get("bf16") is True, "bf16 must be enabled"

    # Verify save_total_limit is small (full checkpoints are ~70GB each)
    save_limit = config.get("save_total_limit", 10)
    assert save_limit <= 5, (
        f"save_total_limit={save_limit} is too high for full SFT. "
        f"Each checkpoint is ~70GB. Limit to 3-5 to avoid filling disk."
    )

    print(f"        Config: {found}")
    print(f"        Optimizer: {optimizer}")
    print(f"        DeepSpeed: {ds_config}")
    print(f"        Learning rate: {lr}")
    print(f"        Model revision: {revision[:12]}...")
    print(f"        bf16: {config.get('bf16')}")
    print(f"        save_total_limit: {save_limit}")
    print(f"        Method: Full-parameter SFT (no LoRA)")


# ── Test 4: DeepSpeed ZeRO-2 config exists and is valid ──────────────────

@test("DeepSpeed ZeRO-2 config is valid")
def test_deepspeed_config():
    ds_paths = [
        "/workspace/runpod-template/ds_config_zero2.json",
        os.path.join(os.path.dirname(__file__), "ds_config_zero2.json"),
    ]
    found = None
    for p in ds_paths:
        if os.path.exists(p):
            found = p
            break

    assert found is not None, (
        f"DeepSpeed config not found. Looked in: {ds_paths}"
    )

    with open(found) as f:
        ds_config = json.load(f)

    # Verify it's ZeRO Stage 2 (not 3)
    stage = ds_config.get("zero_optimization", {}).get("stage")
    assert stage == 2, (
        f"DeepSpeed must be ZeRO Stage 2, got stage {stage}. "
        f"Stage 2 shards gradients; Stage 3 shards params too (not needed for single GPU "
        f"where the model fits in VRAM)."
    )

    # Verify optimizer offload to CPU
    offload_opt = ds_config.get("zero_optimization", {}).get("offload_optimizer", {})
    assert offload_opt.get("device") == "cpu", (
        f"Optimizer must be offloaded to CPU. "
        f"Adafactor states (~35GB) need to live in system RAM."
    )

    # Verify params stay on GPU (not offloaded)
    offload_param = ds_config.get("zero_optimization", {}).get("offload_param", {})
    param_device = offload_param.get("device", "none")
    assert param_device == "none", (
        f"Parameters should NOT be offloaded (device={param_device}). "
        f"The model fits in GPU VRAM — offloading params to CPU would be slow."
    )

    # Verify NO optimizer configured in DeepSpeed (we handle Adafactor in the script)
    assert "optimizer" not in ds_config, (
        "DeepSpeed config should NOT have an optimizer section. "
        "Adafactor is configured in the training script directly — "
        "DeepSpeed's optimizer config conflicts with custom optimizers."
    )

    # Verify bf16 enabled
    assert ds_config.get("bf16", {}).get("enabled") is True, "bf16 must be enabled in DeepSpeed config"

    print(f"        Config: {found}")
    print(f"        ZeRO Stage: {stage}")
    print(f"        Optimizer offload: CPU (pin_memory={offload_opt.get('pin_memory')})")
    print(f"        Param offload: none (stays on GPU)")
    print(f"        bf16: enabled")
    print(f"        No DeepSpeed optimizer block (Adafactor managed by script)")


# ── Test 5: Model architecture loads ─────────────────────────────────────

@test("Model architecture loads")
def test_model_loads():
    import yaml
    from transformers import AutoConfig, AutoTokenizer

    # Read revision from config
    config_paths = [
        "/workspace/runpod-template/train_daimon_config.yaml",
        os.path.join(os.path.dirname(__file__), "train_daimon_config.yaml"),
    ]
    revision = None
    for p in config_paths:
        if os.path.exists(p):
            with open(p) as f:
                cfg = yaml.safe_load(f)
            revision = cfg.get("model_revision")
            break

    model_id = "Qwen/Qwen3.6-35B-A3B"
    local_path = "/workspace/models/Qwen3.6-35B-A3B"

    if os.path.isdir(local_path) and os.path.exists(f"{local_path}/config.json"):
        source = local_path
    else:
        source = model_id

    kwargs = {"trust_remote_code": True}
    if revision and source == model_id:
        kwargs["revision"] = revision

    config = AutoConfig.from_pretrained(source, **kwargs)
    print(f"        Model: {source}")
    print(f"        Type: {config.model_type}")
    print(f"        Hidden: {config.hidden_size}")
    print(f"        Layers: {config.num_hidden_layers}")
    print(f"        Experts: {getattr(config, 'num_experts', 'N/A')}")

    # Also verify tokenizer loads
    tokenizer = AutoTokenizer.from_pretrained(source, **kwargs)
    print(f"        Vocab: {tokenizer.vocab_size}")


# ── Test 6: Training data loads and sequences are within bounds ────────────

@test("Training data loads with valid sequence lengths")
def test_data_loads():
    import yaml
    from datasets import load_from_disk
    from transformers import AutoTokenizer

    data_dir = "/workspace/daimon-data"
    train_arrow = f"{data_dir}/train_arrow"

    assert os.path.isdir(train_arrow), (
        f"Training data not found at {train_arrow}. "
        f"Run setup.sh first to download and prepare data."
    )

    train_ds = load_from_disk(train_arrow)
    print(f"        Train samples: {len(train_ds):,}")

    # Check a sample
    sample = train_ds[0]
    assert "messages" in sample, f"Expected 'messages' key, got: {list(sample.keys())}"
    assert len(sample["messages"]) >= 2, "Each sample needs at least 2 messages (user + assistant)"

    # Verify sequence lengths against max_seq_length
    model_id = "Qwen/Qwen3.6-35B-A3B"
    local_path = "/workspace/models/Qwen3.6-35B-A3B"
    source = local_path if os.path.isdir(local_path) else model_id

    # Read revision from config
    config_paths = [
        "/workspace/runpod-template/train_daimon_config.yaml",
        os.path.join(os.path.dirname(__file__), "train_daimon_config.yaml"),
    ]
    kwargs = {"trust_remote_code": True}
    for p in config_paths:
        if os.path.exists(p):
            with open(p) as f:
                cfg = yaml.safe_load(f)
            revision = cfg.get("model_revision")
            if revision and source == model_id:
                kwargs["revision"] = revision
            break

    tokenizer = AutoTokenizer.from_pretrained(source, **kwargs)

    max_seq_length = 4096  # From config (reduced for full SFT)
    too_long = 0
    max_found = 0

    for i, example in enumerate(train_ds):
        try:
            text = tokenizer.apply_chat_template(
                example["messages"], tokenize=False, add_generation_prompt=False
            )
            tokens = len(tokenizer.encode(text, add_special_tokens=False))
            max_found = max(max_found, tokens)
            if tokens > max_seq_length:
                too_long += 1
        except Exception:
            pass

        if i >= 100:  # Check first 100 samples
            break

    print(f"        Max tokens in sample: {max_found}")
    print(f"        Exceeding {max_seq_length}: {too_long}/{min(len(train_ds), 101)}")

    if too_long > 0:
        print(f"        WARNING: {too_long} sequences exceed max_seq_length.")
        print(f"        The training script will pre-split these, but check your data.")


# ── Test 7: Persistent volume is mounted and writable ──────────────────────

@test("Persistent volume is mounted and writable")
def test_persistent_volume():
    workspace = "/workspace"
    assert os.path.isdir(workspace), "/workspace not found. Is the persistent volume mounted?"

    # Check it's writable
    test_file = os.path.join(workspace, ".daimon_write_test")
    try:
        with open(test_file, "w") as f:
            f.write("test")
        os.remove(test_file)
    except PermissionError:
        raise AssertionError("/workspace is not writable. Check volume permissions.")

    # Check available space — full checkpoints are ~70GB each
    import shutil
    total, used, free = shutil.disk_usage(workspace)
    free_gb = free / (1024**3)
    total_gb = total / (1024**3)
    print(f"        Volume: {total_gb:.0f} GB total, {free_gb:.0f} GB free")

    assert free_gb >= 200, (
        f"Only {free_gb:.0f} GB free on /workspace. "
        f"Need at least 200GB for model + full checkpoints (~70GB each, limit=3)."
    )


# ── Test 8: Memory estimate for full SFT ────────────────────────────────

@test("Memory estimate: full SFT fits in GPU + CPU")
def test_memory_estimate():
    """
    Estimate memory usage for full-parameter SFT with Adafactor + ZeRO-2.
    Verifies both GPU VRAM and system RAM are sufficient.
    Does NOT load the full model — just calculates from config.
    """
    import torch

    # Qwen3.6-35B-A3B has ~35B total params
    total_params = 35e9

    # GPU memory budget
    model_gb = total_params * 2 / 1e9         # bf16 = 2 bytes per param = ~70GB
    activation_gb = 20.0                        # with gradient checkpointing
    gpu_total = model_gb + activation_gb        # ~90GB

    # CPU memory budget (ZeRO-2 offloaded)
    gradient_gb = total_params * 2 / 1e9       # bf16 gradients = ~70GB
    # Adafactor: factored second moments, roughly 1 state per param in mixed precision
    # Much less than AdamW's 2 fp32 states (280GB)
    adafactor_gb = total_params * 1 / 1e9      # ~35GB (conservative estimate)
    cpu_total = gradient_gb + adafactor_gb      # ~105GB

    # AdamW comparison (for reference)
    adamw_gb = total_params * 4 * 2 / 1e9      # 2 fp32 states = ~280GB

    # Available resources
    gpu_vram_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
    ram_bytes = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")
    system_ram_gb = ram_bytes / 1e9

    print(f"        === GPU Memory ===")
    print(f"        Model params (bf16):     {model_gb:.1f} GB")
    print(f"        Activations (grad ckpt): {activation_gb:.1f} GB")
    print(f"        GPU total:               {gpu_total:.1f} GB")
    print(f"        GPU available:           {gpu_vram_gb:.1f} GB")
    print(f"        GPU headroom:            {gpu_vram_gb - gpu_total:.1f} GB")
    print(f"        ")
    print(f"        === CPU Memory (offloaded) ===")
    print(f"        Gradients (bf16):        {gradient_gb:.1f} GB")
    print(f"        Adafactor states:        {adafactor_gb:.1f} GB")
    print(f"        CPU total:               {cpu_total:.1f} GB")
    print(f"        System RAM:              {system_ram_gb:.0f} GB")
    print(f"        CPU headroom:            {system_ram_gb - cpu_total:.0f} GB")
    print(f"        ")
    print(f"        === Why NOT AdamW ===")
    print(f"        AdamW states would need: {adamw_gb:.0f} GB CPU RAM")
    print(f"        System RAM available:    {system_ram_gb:.0f} GB")
    print(f"        Deficit:                 {adamw_gb - system_ram_gb:.0f} GB (does not fit)")

    assert gpu_total < gpu_vram_gb, (
        f"Estimated GPU usage ({gpu_total:.1f} GB) exceeds GPU capacity ({gpu_vram_gb:.1f} GB)."
    )

    assert cpu_total < system_ram_gb * 0.85, (
        f"Estimated CPU usage ({cpu_total:.1f} GB) exceeds safe threshold "
        f"({system_ram_gb * 0.85:.0f} GB = 85% of {system_ram_gb:.0f} GB). "
        f"Need headroom for OS, data loading, and PyTorch buffers."
    )


# ── Test 9: Smoke test — imports and config validation ────────────────────

@test("Training imports and config validation (smoke test)")
def test_smoke():
    """
    Verify all training imports work and the config is valid.
    Does NOT load the full model — that would require too much VRAM for a test.
    """
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, Adafactor

    # Verify trainer imports (no peft needed)
    print(f"        Verifying trainer imports...")
    from trl import SFTTrainer, SFTConfig
    import deepspeed

    # Verify Adafactor is importable
    print(f"        Adafactor: importable from transformers")

    # Verify NO peft dependency
    # (peft may be installed but should not be required)
    print(f"        No peft/LoRA dependency required for full SFT")

    # Create a minimal SFTConfig to verify all parameters are accepted
    test_config = SFTConfig(
        output_dir="/tmp/daimon_test",
        max_length=256,
        num_train_epochs=1,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=1,
        learning_rate=5e-6,
        max_steps=1,
        bf16=True,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        report_to="none",
    )

    print(f"        SFTConfig created successfully")
    print(f"        DeepSpeed version: {deepspeed.__version__}")
    print(f"        TRL version: {__import__('trl').__version__}")
    print(f"        Transformers version: {__import__('transformers').__version__}")

    # Verify YAML config loads
    import yaml
    config_paths = [
        "/workspace/runpod-template/train_daimon_config.yaml",
        os.path.join(os.path.dirname(__file__), "train_daimon_config.yaml"),
    ]
    for p in config_paths:
        if os.path.exists(p):
            with open(p) as f:
                cfg = yaml.safe_load(f)
            print(f"        YAML config loaded: {len(cfg)} keys")
            break

    # Clean up
    import shutil
    if os.path.exists("/tmp/daimon_test"):
        shutil.rmtree("/tmp/daimon_test")

    print(f"        Smoke test passed.")


# ── Run all tests ──────────────────────────────────────────────────────────

def main():
    print("=" * 60)
    print("  DAIMON TRAINING TEMPLATE — VALIDATION TESTS")
    print("  Method: Full-Parameter SFT (no LoRA)")
    print(f"  {time.strftime('%Y-%m-%dT%H:%M:%S')}")
    print("=" * 60)
    print()

    # Collect all test functions
    tests = [v for v in globals().values() if callable(v) and getattr(v, '_test', False)]

    for test_fn in tests:
        test_fn()
        print()

    # Summary
    passed = sum(1 for r in RESULTS if r[0] == "PASS")
    failed = sum(1 for r in RESULTS if r[0] == "FAIL")

    print("=" * 60)
    print(f"  RESULTS: {passed} passed, {failed} failed")
    print()

    for status, name, error in RESULTS:
        marker = "PASS" if status == "PASS" else "FAIL"
        print(f"    [{marker}] {name}")
        if error:
            print(f"           {error}")

    print()
    if failed == 0:
        print("  STATUS: ALL TESTS PASSED — READY TO TRAIN")
        print("  Next: bash /workspace/runpod-template/launch.sh")
    else:
        print("  STATUS: FIX FAILURES BEFORE TRAINING")

    print("=" * 60)

    sys.exit(failed)


if __name__ == "__main__":
    main()