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
WORKER_PORT = 12345
# Cluster spec will be built dynamically based on actual worker hostnames

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 = {
        'worker_addresses': {},  # {task_index: 'hostname:port'}
        'ready_workers': [],
        'completed_workers': [],  # Track workers that finished training
        'training_started': False,
        'start_time': None
    }
    Variable.set('tf_sync_state', json.dumps(sync_state))
    print("Training preparation complete. Sync state initialized.")
    return True


def run_training_node_func(task_index):
    """Execute TensorFlow distributed training"""
    import os
    import socket
    from datetime import datetime
    import time
    import sys
    import tempfile

    # Set NCCL environment variables BEFORE importing TensorFlow
    nccl_log_file = tempfile.mktemp(suffix='.nccl.log')
    os.environ['NCCL_DEBUG'] = 'TRACE'
    os.environ['NCCL_DEBUG_FILE'] = nccl_log_file
    os.environ['NCCL_DEBUG_SUBSYS'] = 'NET'
    os.environ['NCCL_IB_DISABLE'] = '0'  # Enable InfiniBand
    os.environ['NCCL_IB_GPU_DIRECT_RDMA_MAX_DISTANCE'] = '6'
    os.environ['NCCL_P2P_DISABLE'] = '0'
    os.environ['NCCL_IGNORE_CPU_AFFINITY'] = '1'

    hostname = socket.gethostname()
    my_address = f"{hostname}.airflow-worker-gpu:{WORKER_PORT}"

    print(f"{'='*60}")
    print(f"Worker {task_index}/{WORLD_SIZE} - Starting at {datetime.now()}")
    print(f"Hostname: {hostname}")
    print(f"My address: {my_address}")
    print(f"{'='*60}")

    # STEP 1: Register worker address and wait for all workers
    print(f"[Worker-{task_index}] Registering address and waiting for all workers...")
    max_wait = 300  # 5 minutes
    start_wait = time.time()
    cluster_spec = None

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

            # Register my address
            if str(task_index) not in worker_addresses:
                worker_addresses[str(task_index)] = my_address
                sync_state['worker_addresses'] = worker_addresses
                Variable.set('tf_sync_state', json.dumps(sync_state))
                print(f"[Worker-{task_index}] Registered address: {my_address}")

            # Mark as ready
            if task_index not in sync_state.get('ready_workers', []):
                sync_state.setdefault('ready_workers', []).append(task_index)
                Variable.set('tf_sync_state', json.dumps(sync_state))

            # STEP 2: Wait for all workers to register
            if len(worker_addresses) == WORLD_SIZE:
                print(f"[Worker-{task_index}] All {WORLD_SIZE} workers registered!")
                # Build cluster spec from registered addresses
                worker_list = [worker_addresses[str(i)] for i in range(WORLD_SIZE)]
                cluster_spec = {'worker': worker_list}
                print(f"[Worker-{task_index}] Cluster spec: {cluster_spec}")
                break

            print(f"[Worker-{task_index}] Waiting for workers... ({len(worker_addresses)}/{WORLD_SIZE} registered)")
            time.sleep(2)

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

    if cluster_spec is None:
        raise RuntimeError(f"[Worker-{task_index}] Failed to build cluster spec!")

    # Small delay to ensure all workers see the complete cluster
    time.sleep(3)

    # STEP 3: Configure TF_CONFIG BEFORE importing TensorFlow (CRITICAL!)
    tf_config = {
        'cluster': cluster_spec,
        'task': {
            'type': 'worker',
            'index': task_index
        }
    }

    os.environ['TF_CONFIG'] = json.dumps(tf_config)

    print(f"[Worker-{task_index}] TF_CONFIG set:")
    print(f"  Cluster: {cluster_spec}")
    print(f"  Task type: worker")
    print(f"  Task index: {task_index}")
    print(f"  NCCL_IB_DISABLE: {os.environ.get('NCCL_IB_DISABLE')} (IB enabled)")

    # Pre-flight IB diagnostics
    print(f"[Worker-{task_index}] Running IB pre-flight checks...")
    import subprocess
    ib_dev_path = '/dev/infiniband'
    if os.path.exists(ib_dev_path):
        print(f"[Worker-{task_index}] /dev/infiniband: {os.listdir(ib_dev_path)}")
    else:
        print(f"[Worker-{task_index}] /dev/infiniband does NOT exist")
    ib_sys_path = '/sys/class/infiniband'
    if os.path.exists(ib_sys_path):
        ib_devs = os.listdir(ib_sys_path)
        print(f"[Worker-{task_index}] /sys/class/infiniband: {ib_devs}")
        for dev in ib_devs:
            try:
                result = subprocess.run(['ibv_devinfo', '-d', dev], capture_output=True, text=True, timeout=10)
                for line in result.stdout.split('\n'):
                    if any(k in line for k in ['hca_id', 'state', 'link_layer', 'port_lid']):
                        print(f"  {line.strip()}")
            except Exception:
                pass
    else:
        print(f"[Worker-{task_index}] /sys/class/infiniband does NOT exist")

    # NOW import TensorFlow after TF_CONFIG is set
    import tensorflow as tf
    print(f"TensorFlow version: {tf.__version__}")

    # STEP 4: Initialize distribution strategy IMMEDIATELY (CRITICAL!)
    # Must happen before any other TensorFlow operations
    print(f"[Worker-{task_index}] Initializing MultiWorkerMirroredStrategy...")

    # Configure communication options for NCCL
    communication_options = tf.distribute.experimental.CommunicationOptions(
        implementation=tf.distribute.experimental.CommunicationImplementation.NCCL
    )

    strategy = tf.distribute.MultiWorkerMirroredStrategy(
        communication_options=communication_options
    )

    print(f"[Worker-{task_index}] ✓ Strategy initialized!")
    print(f"  Number of devices in sync: {strategy.num_replicas_in_sync}")

    # STEP 5: GPU info (after strategy is initialized)
    gpus = tf.config.list_physical_devices('GPU')
    print(f"[Worker-{task_index}] Available GPUs: {len(gpus)}")
    if gpus:
        print(f"[Worker-{task_index}] GPU: {gpus[0].name}")

    # STEP 6: Define model and training within strategy scope
    with strategy.scope():
        # Create model
        model = tf.keras.Sequential([
            tf.keras.layers.Dense(128, activation='relu', input_shape=(10,)),
            tf.keras.layers.Dense(128, activation='relu'),
            tf.keras.layers.Dense(10)
        ])

        # Compile model
        model.compile(
            optimizer=tf.keras.optimizers.SGD(learning_rate=0.001, momentum=0.9),
            loss=tf.keras.losses.MeanSquaredError(),
            metrics=['mae']
        )

        print(f"[Worker-{task_index}] Model created with {model.count_params()} parameters")

    # STEP 7: Create synthetic dataset with manual sharding
    print(f"[Worker-{task_index}] Creating dataset...")

    import numpy as np

    batch_size = 32
    total_samples = 160
    steps_per_epoch = 5
    num_workers = strategy.num_replicas_in_sync

    # Create synthetic data - same for all workers initially
    np.random.seed(42)
    x_data = np.random.randn(total_samples, 10).astype(np.float32)
    y_data = np.random.randn(total_samples, 10).astype(np.float32)

    # Manually shard data for this worker
    # Each worker gets a disjoint subset
    samples_per_worker = total_samples // num_workers
    start_idx = task_index * samples_per_worker
    end_idx = start_idx + samples_per_worker

    x_worker = x_data[start_idx:end_idx]
    y_worker = y_data[start_idx:end_idx]

    print(f"[Worker-{task_index}] Worker data: indices {start_idx}:{end_idx} ({len(x_worker)} samples)")

    # Create dataset for this worker
    dataset = tf.data.Dataset.from_tensor_slices((x_worker, y_worker))
    dataset = dataset.shuffle(len(x_worker))
    dataset = dataset.batch(batch_size, drop_remainder=True)
    dataset = dataset.repeat()

    # Disable auto-sharding since we manually sharded
    options = tf.data.Options()
    options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
    dataset = dataset.with_options(options)

    print(f"[Worker-{task_index}] Dataset created (batch size: {batch_size})")

    # STEP 8: Custom training loop (avoids PerReplica issues)
    print(f"\n[Worker-{task_index}] {'='*60}")
    print(f"[Worker-{task_index}] Starting Training Loop")
    print(f"[Worker-{task_index}] {'='*60}")

    num_epochs = 10
    num_batches = 2  # Each worker will train on 2 batches per epoch

    @tf.function
    def train_step(x, y):
        with tf.GradientTape() as tape:
            predictions = model(x, training=True)
            loss = model.compiled_loss(y, predictions)

        gradients = tape.gradient(loss, model.trainable_variables)
        model.optimizer.apply_gradients(zip(gradients, model.trainable_variables))

        return loss

    history_loss = []
    training_start = time.time()

    for epoch in range(num_epochs):
        epoch_losses = []
        epoch_start = time.time()

        # Get batches from dataset
        batch_iter = iter(dataset)
        for batch_idx in range(num_batches):
            x_batch, y_batch = next(batch_iter)

            # Distributed training step
            per_replica_loss = strategy.run(train_step, args=(x_batch, y_batch))

            # Reduce loss across replicas
            loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_loss, axis=None)
            epoch_losses.append(float(loss.numpy()))

        avg_loss = sum(epoch_losses) / len(epoch_losses)
        history_loss.append(avg_loss)
        epoch_time = time.time() - epoch_start

        print(f"[Worker-{task_index}] Epoch {epoch+1}/{num_epochs} DONE | Loss: {avg_loss:.6f} | Time: {epoch_time:.3f}s")

    total_time = time.time() - training_start

    print(f"\n[Worker-{task_index}] {'='*60}")
    print(f"[Worker-{task_index}] Training Complete!")
    print(f"[Worker-{task_index}] Total training time: {total_time:.3f}s")
    print(f"[Worker-{task_index}] {'='*60}")

    # Print NCCL transport info from log file
    print(f"\n[Worker-{task_index}] === NCCL Transport Summary ===")
    try:
        with open(nccl_log_file, 'r') as f:
            nccl_lines = f.readlines()
        for line in nccl_lines:
            line = line.strip()
            if any(kw in line for kw in ['NET/IB', 'GPU Direct', 'Channel', 'via NET', 'via Socket', 'transport']):
                print(f"[Worker-{task_index}] NCCL: {line}")
    except Exception as e:
        print(f"[Worker-{task_index}] Could not read NCCL log: {e}")

    # Get final metrics
    final_loss = history_loss[-1]

    if task_index == 0:
        print(f"[Worker-{task_index}] Final Loss: {final_loss:.6f}")

    # STEP 9: Synchronization barrier - wait for all workers to finish training
    print(f"[Worker-{task_index}] Waiting at barrier for all workers to complete...")
    import time
    max_wait = 120
    start_wait = time.time()

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

            if task_index not in completed_workers:
                completed_workers.append(task_index)
                sync_state['completed_workers'] = completed_workers
                Variable.set('tf_sync_state', json.dumps(sync_state))
                print(f"[Worker-{task_index}] Marked as completed ({len(completed_workers)}/{WORLD_SIZE})")

            if len(completed_workers) == WORLD_SIZE:
                print(f"[Worker-{task_index}] All workers completed! Proceeding to cleanup.")
                break

            time.sleep(2)
        except Exception as e:
            print(f"[Worker-{task_index}] Error during completion sync: {e}")
            time.sleep(2)
    else:
        print(f"[Worker-{task_index}] Warning: Timeout at completion barrier")

    time.sleep(2)  # Small delay to ensure all workers pass barrier

    # STEP 10: Save model (only chief worker)
    if task_index == 0:
        model_path = '/tmp/tf_distributed_model.keras'
        print(f"[Worker-{task_index}] Saving model to {model_path}...")
        model.save(model_path)
        print(f"[Worker-{task_index}] ✓ Model saved successfully!")

    print(f"[Worker-{task_index}] Exiting gracefully...")

    return {
        'worker_index': task_index,
        'final_loss': float(final_loss),
        'epochs_completed': num_epochs,
        'status': 'success'
    }


def cleanup_sync_state_func():
    """Clean up synchronization state"""
    try:
        Variable.delete('tf_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"TENSORFLOW DISTRIBUTED TRAINING SUMMARY")
    print(f"{'='*60}")

    for worker_idx in range(WORLD_SIZE):
        try:
            result = ti.xcom_pull(task_ids=f"train_worker_{worker_idx}")
            if result and result.get('status') == 'success':
                print(f"Worker {result['worker_index']}: ✓ 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 worker {worker_idx}: {e}")

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


with DAG(
    dag_id='tensorflow_distributed_training_multiworker',
    default_args=default_args,
    schedule=None,
    start_date=pendulum.today('UTC').add(days=-1),
    catchup=False,
    tags=['gpu', 'ml', 'distributed', 'tensorflow'],
    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 worker
    training_tasks = []
    for i in range(WORLD_SIZE):
        task = PythonOperator(
            task_id=f'train_worker_{i}',
            python_callable=run_training_node_func,
            op_kwargs={'task_index': i},
            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