//! HPC Channels integration for distributed training.
//!
//! This module bridges distributed training events to the hpc-channels message bus,
//! enabling real-time monitoring and coordination of multi-GPU/multi-node training.
//!
//! # Channels Used
//!
//! - `hpc.training.start` - Training job started
//! - `hpc.training.progress` - Training progress updates
//! - `hpc.training.metrics` - Training metrics (loss, accuracy, throughput)
//! - `hpc.training.status` - Training status changes
//! - `hpc.training.control` - Training control signals (pause, resume, stop)
//!
//! # Example
//!
//! ```rust,ignore
//! use rtx_distributed::channels::TrainingChannelBridge;
//!
//! let bridge = TrainingChannelBridge::new();
//!
//! // Publish training started
//! bridge.publish_training_start(&config);
//!
//! // Publish progress update
//! bridge.publish_progress(epoch, step, total_steps);
//!
//! // Publish metrics
//! bridge.publish_metrics(&metrics);
//! ```

use std::sync::Arc;

use tokio::sync::broadcast;

/// Training job started event.
#[derive(Clone, Debug)]
pub struct TrainingStartEvent {
    /// Job identifier.
    pub job_id: String,
    /// Model name.
    pub model_name: String,
    /// Total number of epochs.
    pub total_epochs: u32,
    /// Batch size per GPU.
    pub batch_size: u32,
    /// Number of GPUs.
    pub num_gpus: u32,
    /// World size (total processes).
    pub world_size: u32,
    /// Timestamp (epoch ms).
    pub timestamp_ms: u64,
}

/// Training progress update.
#[derive(Clone, Debug)]
pub struct TrainingProgressEvent {
    /// Job identifier.
    pub job_id: String,
    /// Current epoch.
    pub epoch: u32,
    /// Current step within epoch.
    pub step: u64,
    /// Total steps in epoch.
    pub total_steps: u64,
    /// Estimated time remaining (seconds).
    pub eta_seconds: Option<u64>,
    /// Timestamp (epoch ms).
    pub timestamp_ms: u64,
}

/// Training metrics snapshot.
#[derive(Clone, Debug)]
pub struct TrainingMetricsEvent {
    /// Job identifier.
    pub job_id: String,
    /// Current loss.
    pub loss: f64,
    /// Learning rate.
    pub learning_rate: f64,
    /// Throughput (samples/second).
    pub throughput: f64,
    /// GPU memory used (bytes).
    pub gpu_memory_used: u64,
    /// GPU utilization (0-100).
    pub gpu_utilization: f32,
    /// Gradient norm.
    pub grad_norm: Option<f64>,
    /// Timestamp (epoch ms).
    pub timestamp_ms: u64,
}

/// Training status.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum TrainingStatus {
    /// Training is initializing.
    Initializing,
    /// Training is running.
    Running,
    /// Training is paused.
    Paused,
    /// Training completed successfully.
    Completed,
    /// Training failed.
    Failed(String),
    /// Training was cancelled.
    Cancelled,
}

/// Training status change event.
#[derive(Clone, Debug)]
pub struct TrainingStatusEvent {
    /// Job identifier.
    pub job_id: String,
    /// New status.
    pub status: TrainingStatus,
    /// Timestamp (epoch ms).
    pub timestamp_ms: u64,
}

/// Training control commands.
#[derive(Clone, Debug)]
pub enum TrainingControlCommand {
    /// Pause training.
    Pause,
    /// Resume training.
    Resume,
    /// Stop training gracefully.
    Stop,
    /// Checkpoint now.
    CheckpointNow,
    /// Adjust learning rate.
    AdjustLr(f64),
}

/// Bridge between distributed training and hpc-channels.
pub struct TrainingChannelBridge {
    /// Broadcast sender for training start events.
    start_tx: broadcast::Sender<TrainingStartEvent>,
    /// Broadcast sender for progress events.
    progress_tx: broadcast::Sender<TrainingProgressEvent>,
    /// Broadcast sender for metrics events.
    metrics_tx: broadcast::Sender<TrainingMetricsEvent>,
    /// Broadcast sender for status events.
    status_tx: broadcast::Sender<TrainingStatusEvent>,
    /// Broadcast sender for control commands.
    control_tx: broadcast::Sender<TrainingControlCommand>,
}

impl TrainingChannelBridge {
    /// Create a new training channel bridge.
    ///
    /// Registers channels with the hpc-channels global registry.
    pub fn new() -> Self {
        let start_tx = hpc_channels::broadcast::<TrainingStartEvent>(
            hpc_channels::channels::TRAINING_START,
            256,
        );
        let progress_tx = hpc_channels::broadcast::<TrainingProgressEvent>(
            hpc_channels::channels::TRAINING_PROGRESS,
            1024,
        );
        let metrics_tx = hpc_channels::broadcast::<TrainingMetricsEvent>(
            hpc_channels::channels::TRAINING_METRICS,
            1024,
        );
        let status_tx = hpc_channels::broadcast::<TrainingStatusEvent>(
            hpc_channels::channels::TRAINING_STATUS,
            256,
        );
        let control_tx = hpc_channels::broadcast::<TrainingControlCommand>(
            hpc_channels::channels::TRAINING_CONTROL,
            64,
        );

        Self {
            start_tx,
            progress_tx,
            metrics_tx,
            status_tx,
            control_tx,
        }
    }

    fn now_ms() -> u64 {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_millis() as u64)
            .unwrap_or(0)
    }

    /// Publish training started event.
    pub fn publish_training_start(
        &self,
        job_id: &str,
        model_name: &str,
        total_epochs: u32,
        batch_size: u32,
        num_gpus: u32,
        world_size: u32,
    ) {
        let _ = self.start_tx.send(TrainingStartEvent {
            job_id: job_id.to_string(),
            model_name: model_name.to_string(),
            total_epochs,
            batch_size,
            num_gpus,
            world_size,
            timestamp_ms: Self::now_ms(),
        });
    }

    /// Publish training progress.
    pub fn publish_progress(
        &self,
        job_id: &str,
        epoch: u32,
        step: u64,
        total_steps: u64,
        eta_seconds: Option<u64>,
    ) {
        let _ = self.progress_tx.send(TrainingProgressEvent {
            job_id: job_id.to_string(),
            epoch,
            step,
            total_steps,
            eta_seconds,
            timestamp_ms: Self::now_ms(),
        });
    }

    /// Publish training metrics.
    pub fn publish_metrics(
        &self,
        job_id: &str,
        loss: f64,
        learning_rate: f64,
        throughput: f64,
        gpu_memory_used: u64,
        gpu_utilization: f32,
        grad_norm: Option<f64>,
    ) {
        let _ = self.metrics_tx.send(TrainingMetricsEvent {
            job_id: job_id.to_string(),
            loss,
            learning_rate,
            throughput,
            gpu_memory_used,
            gpu_utilization,
            grad_norm,
            timestamp_ms: Self::now_ms(),
        });
    }

    /// Publish training status change.
    pub fn publish_status(&self, job_id: &str, status: TrainingStatus) {
        let _ = self.status_tx.send(TrainingStatusEvent {
            job_id: job_id.to_string(),
            status,
            timestamp_ms: Self::now_ms(),
        });
    }

    /// Send a control command.
    pub fn send_control(&self, command: TrainingControlCommand) {
        let _ = self.control_tx.send(command);
    }

    /// Subscribe to training start events.
    pub fn subscribe_start(&self) -> broadcast::Receiver<TrainingStartEvent> {
        self.start_tx.subscribe()
    }

    /// Subscribe to progress events.
    pub fn subscribe_progress(&self) -> broadcast::Receiver<TrainingProgressEvent> {
        self.progress_tx.subscribe()
    }

    /// Subscribe to metrics events.
    pub fn subscribe_metrics(&self) -> broadcast::Receiver<TrainingMetricsEvent> {
        self.metrics_tx.subscribe()
    }

    /// Subscribe to status events.
    pub fn subscribe_status(&self) -> broadcast::Receiver<TrainingStatusEvent> {
        self.status_tx.subscribe()
    }

    /// Subscribe to control commands.
    pub fn subscribe_control(&self) -> broadcast::Receiver<TrainingControlCommand> {
        self.control_tx.subscribe()
    }
}

impl Default for TrainingChannelBridge {
    fn default() -> Self {
        Self::new()
    }
}

/// Shared channel bridge type.
pub type SharedTrainingChannelBridge = Arc<TrainingChannelBridge>;

/// Create a new shared channel bridge.
#[must_use]
pub fn shared_channel_bridge() -> SharedTrainingChannelBridge {
    Arc::new(TrainingChannelBridge::new())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_bridge_creation() {
        let bridge = TrainingChannelBridge::new();
        assert!(hpc_channels::exists(hpc_channels::channels::TRAINING_START));
        assert!(hpc_channels::exists(
            hpc_channels::channels::TRAINING_PROGRESS
        ));
        assert!(hpc_channels::exists(
            hpc_channels::channels::TRAINING_METRICS
        ));
        let _ = bridge;
    }

    #[tokio::test]
    async fn test_training_start_publishing() {
        let bridge = TrainingChannelBridge::new();
        let mut rx = bridge.subscribe_start();

        bridge.publish_training_start("job-123", "resnet50", 100, 32, 8, 8);

        let event = rx.recv().await.expect("Should receive event");
        assert_eq!(event.job_id, "job-123");
        assert_eq!(event.model_name, "resnet50");
        assert_eq!(event.total_epochs, 100);
        assert_eq!(event.batch_size, 32);
    }

    #[tokio::test]
    async fn test_progress_publishing() {
        let bridge = TrainingChannelBridge::new();
        let mut rx = bridge.subscribe_progress();

        bridge.publish_progress("job-123", 5, 1000, 5000, Some(3600));

        let event = rx.recv().await.expect("Should receive event");
        assert_eq!(event.epoch, 5);
        assert_eq!(event.step, 1000);
        assert_eq!(event.total_steps, 5000);
    }

    #[tokio::test]
    async fn test_metrics_publishing() {
        let bridge = TrainingChannelBridge::new();
        let mut rx = bridge.subscribe_metrics();

        bridge.publish_metrics(
            "job-123",
            0.05,      // loss
            0.001,     // learning rate
            1500.0,    // throughput
            8_000_000, // gpu memory
            95.0,      // gpu utilization
            Some(1.5), // grad norm
        );

        let event = rx.recv().await.expect("Should receive event");
        assert!((event.loss - 0.05).abs() < f64::EPSILON);
        assert!((event.throughput - 1500.0).abs() < f64::EPSILON);
    }

    #[tokio::test]
    async fn test_status_publishing() {
        let bridge = TrainingChannelBridge::new();
        let mut rx = bridge.subscribe_status();

        bridge.publish_status("job-123", TrainingStatus::Running);

        let event = rx.recv().await.expect("Should receive event");
        assert_eq!(event.status, TrainingStatus::Running);
    }

    #[tokio::test]
    async fn test_control_commands() {
        let bridge = TrainingChannelBridge::new();
        let mut rx = bridge.subscribe_control();

        bridge.send_control(TrainingControlCommand::Pause);

        let cmd = rx.recv().await.expect("Should receive command");
        assert!(matches!(cmd, TrainingControlCommand::Pause));
    }
}