from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.models import Variable
import pendulum
from datetime import timedelta
import json

# --- CONFIGURATION ---
WORLD_SIZE = 2
MASTER_ADDR = "airflow-worker-gpu-0.airflow-worker-gpu.george-tests.svc.cluster.local"
MASTER_PORT = "29500"
NCCL_TIMEOUT = 1800
NAMESPACE = "george-tests"

default_args = {
  'owner': 'airflow',
  'retries': 1,
  'retry_delay': timedelta(minutes=2),
  'execution_timeout': timedelta(hours=2),
}


def prepare_training_func():
  """Initialize shared state for synchronization"""
  sync_state = {
      'ready_workers': [],
      'training_started': False,
      'start_time': None
  }
  Variable.set('ddp_sync_state', json.dumps(sync_state))
  print("Training preparation complete. Sync state initialized.")
  return True


def run_training_node_func(rank, world_size):
  """Execute distributed training with synchronization"""
  import os
  import socket
  from datetime import datetime
  import time
  import sys

  # Redirect NCCL stderr to stdout so Airflow captures it
  os.dup2(sys.stdout.fileno(), sys.stderr.fileno())

  # CRITICAL: Set NCCL environment variables BEFORE importing torch
  os.environ['NCCL_DEBUG'] = 'INFO'
  os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT)
  os.environ['TORCH_NCCL_BLOCKING_WAIT'] = '1'
  os.environ['NCCL_IB_DISABLE'] = '0'  # Enable InfiniBand
  os.environ['NCCL_IB_GPU_DIRECT_RDMA_MAX_DISTANCE'] = '6'  # GPU at bus 0, IB at bus 6
  os.environ['NCCL_P2P_DISABLE'] = '0'  # Enable P2P
  os.environ['NCCL_IGNORE_CPU_AFFINITY'] = '1'  # Better compatibility

  # Now import torch after NCCL config is set
  import torch
  import torch.distributed as dist
  import torch.nn as nn
  import torch.optim as optim

  print(f"{'='*60}")
  print(f"Node Rank {rank}/{world_size} - Starting at {datetime.now()}")
  print(f"Hostname: {socket.gethostname()}")
  print(f"{'='*60}")

  # STEP 1: Signal that this worker is ready
  print(f"[{rank}] Signaling ready state...")
  max_wait = 300  # 5 minutes
  start_wait = time.time()

  while time.time() - start_wait < max_wait:
      try:
          sync_state = json.loads(Variable.get('ddp_sync_state', default_var='{}'))

          if rank not in sync_state.get('ready_workers', []):
              sync_state.setdefault('ready_workers', []).append(rank)
              Variable.set('ddp_sync_state', json.dumps(sync_state))
              print(f"[{rank}] Marked as ready. Ready workers: {sync_state['ready_workers']}")

          # STEP 2: Wait for all workers to be ready
          if len(sync_state.get('ready_workers', [])) == world_size:
              print(f"[{rank}] All {world_size} workers are ready! Proceeding to training...")
              break

          print(f"[{rank}] Waiting for other workers... ({len(sync_state.get('ready_workers', []))}/{world_size} ready)")
          time.sleep(2)

      except Exception as e:
          print(f"[{rank}] Error during sync: {e}")
          time.sleep(2)
  else:
      raise RuntimeError(f"[{rank}] Timeout waiting for all workers to be ready!")

  # STEP 2.5: Rank 0 stores its address, others wait for it
  print(f"[{rank}] Determining master address...")
  my_fqdn = socket.getfqdn()
  print(f"[{rank}] My FQDN: {my_fqdn}")

  master_addr = None  # Initialize outside if/else
  if rank == 0:
      # Rank 0 stores its address as the master
      sync_state = json.loads(Variable.get('ddp_sync_state', default_var='{}'))
      sync_state['master_addr'] = my_fqdn
      Variable.set('ddp_sync_state', json.dumps(sync_state))
      master_addr = my_fqdn  # For rank 0, master is itself
      print(f"[{rank}] I am the master. Stored my address: {my_fqdn}")
  else:
      # Other ranks wait for rank 0 to store its address
      max_wait = 60
      start_wait = time.time()
      while time.time() - start_wait < max_wait:
          sync_state = json.loads(Variable.get('ddp_sync_state', default_var='{}'))
          if 'master_addr' in sync_state:
              master_addr = sync_state['master_addr']
              print(f"[{rank}] Master address retrieved: {master_addr}")
              break
          print(f"[{rank}] Waiting for master address...")
          time.sleep(1)
      else:
          raise RuntimeError(f"[{rank}] Timeout waiting for master address!")

  # Small delay to ensure all workers see the ready state
  time.sleep(3)

  # STEP 3: Configure distributed environment
  # Use the master_addr that was determined above
  print(f"[{rank}] Using master address: {master_addr}")
  os.environ['MASTER_ADDR'] = master_addr
  os.environ['MASTER_PORT'] = MASTER_PORT
  os.environ['WORLD_SIZE'] = str(world_size)
  os.environ['RANK'] = str(rank)

  print(f"[{rank}] Environment configured:")
  print(f"  MASTER_ADDR: {master_addr}")
  print(f"  MASTER_PORT: {MASTER_PORT}")
  print(f"  RANK: {rank}")
  print(f"  WORLD_SIZE: {world_size}")
  print(f"  NCCL_IB_DISABLE: {os.environ.get('NCCL_IB_DISABLE')} (IB enabled)")

  # STEP 3.5: Pre-flight checks
  print(f"[{rank}] Running pre-flight checks...")
  print(f"[{rank}] Hostname: {socket.gethostname()}")
  print(f"[{rank}] FQDN: {socket.getfqdn()}")

  # Test DNS resolution
  try:
      import subprocess
      result = subprocess.run(['getent', 'hosts', MASTER_ADDR], capture_output=True, text=True)
      print(f"[{rank}] DNS resolution for {MASTER_ADDR}:")
      print(f"  {result.stdout.strip()}")
  except Exception as e:
      print(f"[{rank}] Warning: Could not verify DNS: {e}")

  # Check network interfaces
  try:
      result = subprocess.run(['ip', 'addr'], capture_output=True, text=True)
      ib_lines = [line for line in result.stdout.split('\n') if 'ibp' in line or 'ib0' in line]
      if ib_lines:
          print(f"[{rank}] Found IB interfaces:")
          for line in ib_lines[:5]:
              print(f"  {line.strip()}")
      else:
          print(f"[{rank}] No IB interfaces found via 'ip addr'")
  except Exception as e:
      print(f"[{rank}] Warning: Could not list interfaces: {e}")

  # Check /dev/infiniband devices
  print(f"[{rank}] Checking IB device access...")
  ib_dev_path = '/dev/infiniband'
  if os.path.exists(ib_dev_path):
      ib_devices = os.listdir(ib_dev_path)
      print(f"[{rank}] /dev/infiniband contents: {ib_devices if ib_devices else 'EMPTY'}")
  else:
      print(f"[{rank}] /dev/infiniband does NOT exist - IB devices not mounted in container")

  # Check /sys/class/infiniband
  ib_sys_path = '/sys/class/infiniband'
  if os.path.exists(ib_sys_path):
      ib_sys_devices = os.listdir(ib_sys_path)
      print(f"[{rank}] /sys/class/infiniband contents: {ib_sys_devices if ib_sys_devices else 'EMPTY'}")
  else:
      print(f"[{rank}] /sys/class/infiniband does NOT exist")

  # Try ibv_devinfo
  try:
      result = subprocess.run(['ibv_devinfo'], capture_output=True, text=True, timeout=10)
      print(f"[{rank}] ibv_devinfo output:")
      for line in result.stdout.split('\n')[:20]:
          print(f"  {line}")
      if result.returncode != 0:
          print(f"[{rank}] ibv_devinfo stderr: {result.stderr.strip()}")
  except FileNotFoundError:
      print(f"[{rank}] ibv_devinfo not found - rdma-core tools not installed")
  except Exception as e:
      print(f"[{rank}] ibv_devinfo failed: {e}")

  # Check RDMA devices via sysfs
  rdma_path = '/sys/class/infiniband_verbs'
  if os.path.exists(rdma_path):
      print(f"[{rank}] RDMA verbs devices: {os.listdir(rdma_path)}")
  else:
      print(f"[{rank}] /sys/class/infiniband_verbs does NOT exist")

  # STEP 4: Setup GPU FIRST (before process group init)
  device = torch.device("cuda:0")
  torch.cuda.set_device(device)

  print(f"[{rank}] GPU Configuration:")
  print(f"  Device: {torch.cuda.get_device_name(0)}")
  print(f"  Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
  print(f"  Current device: {torch.cuda.current_device()}")

  # STEP 5: Initialize process group
  print(f"[{rank}] Initializing process group (backend=nccl)...")
  print(f"[{rank}] This may take 30-60 seconds...")

  # Flush output to ensure logs are visible
  import sys
  sys.stdout.flush()

  try:
      dist.init_process_group(
          backend="nccl",
          init_method="env://",
          timeout=timedelta(minutes=10),
          rank=rank,
          world_size=world_size
      )
      print(f"[{rank}] ✓ Successfully joined distributed group!")
      print(f"  Process group size: {dist.get_world_size()}")
      print(f"  My rank: {dist.get_rank()}")
  except Exception as e:
      print(f"[{rank}] ✗ Failed to initialize process group!")
      print(f"  Error type: {type(e).__name__}")
      print(f"  Error: {str(e)}")
      import traceback
      traceback.print_exc()
      raise

  # STEP 6: Define model
  model = nn.Sequential(
      nn.Linear(10, 128),
      nn.ReLU(),
      nn.Linear(128, 128),
      nn.ReLU(),
      nn.Linear(128, 10)
  ).to(device)

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

  criterion = nn.MSELoss()
  optimizer = optim.SGD(ddp_model.parameters(), lr=0.001, momentum=0.9)

  print(f"[{rank}] Model initialized with {sum(p.numel() for p in model.parameters())} parameters")

  # STEP 7: Training loop
  print(f"\n[{rank}] {'='*60}")
  print(f"[{rank}] Starting Training Loop")
  print(f"[{rank}] {'='*60}")

  num_epochs = 10
  batch_size = 32
  num_batches = 5

  training_start = time.time()

  for epoch in range(num_epochs):
      ddp_model.train()
      epoch_loss = 0.0
      epoch_start = time.time()

      for batch_idx in range(num_batches):
          torch.manual_seed(epoch * num_batches + batch_idx)
          inputs = torch.randn(batch_size, 10).to(device)
          labels = torch.randn(batch_size, 10).to(device)

          optimizer.zero_grad()
          outputs = ddp_model(inputs)
          loss = criterion(outputs, labels)
          loss.backward()
          optimizer.step()

          epoch_loss += loss.item()

      avg_loss = epoch_loss / num_batches
      epoch_time = time.time() - epoch_start

      # Synchronize loss across ranks
      loss_tensor = torch.tensor([avg_loss]).to(device)
      dist.all_reduce(loss_tensor, op=dist.ReduceOp.AVG)
      global_avg_loss = loss_tensor.item()

      print(f"[{rank}] Epoch {epoch+1}/{num_epochs} DONE | Local Avg: {avg_loss:.6f} | Global Avg: {global_avg_loss:.6f} | Time: {epoch_time:.3f}s")

  total_time = time.time() - training_start

  print(f"\n[{rank}] {'='*60}")
  print(f"[{rank}] Training Complete!")
  print(f"[{rank}] Total training time: {total_time:.3f}s")
  print(f"[{rank}] {'='*60}")
  dist.destroy_process_group()
  print(f"[{rank}] Process group destroyed. Finished at {datetime.now()}")

  return {
      'rank': rank,
      'final_loss': global_avg_loss if rank == 0 else avg_loss,
      'epochs_completed': num_epochs,
      'status': 'success'
  }


def cleanup_sync_state_func():
  """Clean up synchronization state"""
  try:
      Variable.delete('ddp_sync_state')
      print("Sync state cleaned up")
  except:
      pass
  return True


def training_summary_func(**context):
  """Aggregate and display training results"""
  ti = context['ti']

  print(f"\n{'='*60}")
  print(f"DISTRIBUTED TRAINING SUMMARY")
  print(f"{'='*60}")

  for rank in range(WORLD_SIZE):
      try:
          result = ti.xcom_pull(task_ids=f"train_rank_{rank}")
          if result and result.get('status') == 'success':
              print(f"Rank {result['rank']}: ✓ Completed {result['epochs_completed']} epochs")
              print(f"  Final loss: {result['final_loss']:.6f}")
      except Exception as e:
          print(f"Warning: Could not get result for rank {rank}: {e}")

  print(f"\n✓ Training completed!")
  return {'status': 'success', 'workers': WORLD_SIZE}


with DAG(
  dag_id='pytorch_distributed_training_ddp_production',
  default_args=default_args,
  schedule=None,
  start_date=pendulum.today('UTC').add(days=-1),
  catchup=False,
  tags=['gpu', 'ml', 'distributed'],
  max_active_runs=1,
  max_active_tasks=10,
) as dag:

  # Preparation task
  prep = PythonOperator(
      task_id='prepare_training',
      python_callable=prepare_training_func,
      queue='gpu'
  )

  # Training tasks - one per rank
  training_tasks = []
  for i in range(WORLD_SIZE):
      task = PythonOperator(
          task_id=f'train_rank_{i}',
          python_callable=run_training_node_func,
          op_kwargs={'rank': i, 'world_size': WORLD_SIZE},
          queue='gpu',
      )
      training_tasks.append(task)

  # Cleanup task
  cleanup = PythonOperator(
      task_id='cleanup_sync_state',
      python_callable=cleanup_sync_state_func,
      queue='gpu',
      trigger_rule='all_done'
  )

  # Summary task
  summary = PythonOperator(
      task_id='training_summary',
      python_callable=training_summary_func,
      queue='gpu',
      trigger_rule='all_done'
  )

  # Set dependencies
  prep >> training_tasks >> cleanup >> summary