from airflow import DAG
from airflow.operators.python import PythonOperator
import pendulum
from datetime import timedelta

WORLD_SIZE = 2

MASTER_ADDR = "airflow-worker-gpu-0.airflow-worker-gpu"
MASTER_PORT = "29500"

default_args = {
    "owner": "airflow",
    "retries": 0,
    "execution_timeout": timedelta(hours=2),
}


# -------------------------
# SIMPLE PREP (NO VARIABLES)
# -------------------------
def prepare_training_func():
    print("Starting distributed training (no sync state needed)")
    return True


# -------------------------
# CORE TRAINING
# -------------------------
def run_training_node_func(rank, world_size):
    import os
    import time
    import torch
    import torch.distributed as dist
    import torch.nn as nn
    import torch.optim as optim
    from datetime import datetime
    import socket

    print("=" * 80)
    print(f"RANK {rank}/{world_size} START {datetime.now()}")
    print(f"HOST: {socket.gethostname()}")
    print("=" * 80)

    # -------------------------
    # FIXED BARRIER (important)
    # -------------------------
    print(f"[{rank}] sync barrier (static sleep)")
    time.sleep(10)

    # -------------------------
    # STATIC CONFIG (NO AIRFLOW VARIABLES)
    # -------------------------
    os.environ["MASTER_ADDR"] = MASTER_ADDR
    os.environ["MASTER_PORT"] = MASTER_PORT
    os.environ["WORLD_SIZE"] = str(world_size)
    os.environ["RANK"] = str(rank)

    os.environ["NCCL_SOCKET_IFNAME"] = "eth0"
    os.environ["NCCL_DEBUG"] = "INFO"
    os.environ["TORCH_NCCL_BLOCKING_WAIT"] = "1"

    print(f"[{rank}] MASTER = {MASTER_ADDR}:{MASTER_PORT}")

    # -------------------------
    # INIT PROCESS GROUP
    # -------------------------
    dist.init_process_group(
        backend="nccl",
        init_method="env://",
        rank=rank,
        world_size=world_size,
    )

    print(f"[{rank}] DDP INIT OK")

    # -------------------------
    # GPU BINDING (CRITICAL FIX)
    # -------------------------
    torch.cuda.set_device(0)
    device = torch.device("cuda:0")

    print(f"[{rank}] GPU = {torch.cuda.get_device_name(0)}")

    # -------------------------
    # MODEL
    # -------------------------
    model = nn.Sequential(
        nn.Linear(10, 128),
        nn.ReLU(),
        nn.Linear(128, 10),
    ).to(device)

    ddp_model = nn.parallel.DistributedDataParallel(
        model,
        device_ids=[0],
    )

    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
    loss_fn = nn.MSELoss()

    # -------------------------
    # TRAIN LOOP
    # -------------------------
    for epoch in range(5):
        x = torch.randn(32, 10).to(device)
        y = torch.randn(32, 10).to(device)

        optimizer.zero_grad()
        out = ddp_model(x)
        loss = loss_fn(out, y)
        loss.backward()
        optimizer.step()

        print(f"[{rank}] epoch={epoch} loss={loss.item():.4f}")

    dist.destroy_process_group()

    return {"rank": rank, "status": "ok"}


# -------------------------
# CLEANUP
# -------------------------
def cleanup_func():
    print("cleanup done")
    return True


def summary_func(**context):
    print("training done")
    return True


# -------------------------
# DAG
# -------------------------
with DAG(
    dag_id="ddp_airflow_stable",
    start_date=pendulum.today("UTC").add(days=-1),
    schedule=None,
    catchup=False,
    max_active_runs=1,
    max_active_tasks=2,
    default_args=default_args,
) as dag:

    prep = PythonOperator(
        task_id="prep",
        python_callable=prepare_training_func,
        queue="gpu",
    )

    tasks = []
    for r in range(WORLD_SIZE):
        t = PythonOperator(
            task_id=f"train_rank_{r}",
            python_callable=run_training_node_func,
            op_kwargs={"rank": r, "world_size": WORLD_SIZE},
            queue="gpu",
        )
        tasks.append(t)

    cleanup = PythonOperator(
        task_id="cleanup",
        python_callable=cleanup_func,
        trigger_rule="all_done",
        queue="gpu",
    )

    summary = PythonOperator(
        task_id="summary",
        python_callable=summary_func,
        trigger_rule="all_done",
        queue="gpu",
    )

    prep >> tasks >> cleanup >> summary