//! gRPC Streaming Progress Implementation
//!
//! This module implements real-time progress streaming for ML training jobs using a hybrid approach:
//! 1. Fetch initial state from orchestrator (immediate snapshot)
//! 2. Subscribe to real-time broadcast updates (low-latency live updates)
//! 3. Support both single job and batch (multiplexed) streaming
//! 4. Delta-only updates to minimize bandwidth

use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;

use tokio::sync::mpsc;
use tokio_stream::{wrappers::ReceiverStream, Stream};
use tonic::{Request, Response, Status};
use tracing::{debug, info, warn};
use uuid::Uuid;

use crate::orchestrator::{JobStatus, TrainingOrchestrator, TrainingStatusUpdate};
use crate::proto::ml_training::{
    StreamProgressRequest, ProgressUpdate, TuningJobStatus, UpdateType,
};

/// Progress streaming handler
pub struct ProgressStreamer {
    orchestrator: Arc<TrainingOrchestrator>,
}

impl ProgressStreamer {
    /// Create a new progress streamer
    pub fn new(orchestrator: Arc<TrainingOrchestrator>) -> Self {
        Self { orchestrator }
    }

    /// Stream training progress for a single job or batch
    ///
    /// This implements the hybrid approach:
    /// 1. Fetch current state from orchestrator (initial snapshot)
    /// 2. Subscribe to broadcast channel for real-time updates
    /// 3. Stream delta-only updates to client
    /// 4. Handle client disconnects gracefully
    pub async fn stream_progress(
        &self,
        request: Request<StreamProgressRequest>,
    ) -> Result<Response<Pin<Box<dyn Stream<Item = Result<ProgressUpdate, Status>> + Send>>>, Status>
    {
        let req = request.into_inner();

        info!(
            "Starting progress stream for job_id: {}",
            req.job_id
        );

        // Parse job ID
        let job_id = Uuid::parse_str(&req.job_id)
            .map_err(|_| Status::invalid_argument("Invalid job ID format"))?;

        // Check if job exists
        let initial_job = self
            .orchestrator
            .get_job(job_id)
            .await
            .map_err(|e| Status::not_found(format!("Job not found: {}", e)))?;

        debug!("Found job {} - model: {}, status: {:?}", 
            job_id, initial_job.model_type, initial_job.status);

        // Create channel for streaming
        let (tx, rx) = mpsc::channel(100);

        // Subscribe to job status broadcasts
        let mut status_receiver = self
            .orchestrator
            .subscribe_to_job_status(job_id)
            .await
            .map_err(|e| Status::internal(format!("Failed to subscribe: {}", e)))?;

        // Spawn task to handle streaming
        let orchestrator = self.orchestrator.clone();
        tokio::spawn(async move {
            // Send initial state first (hybrid approach - step 1)
            let initial_update = Self::create_initial_progress_update(&initial_job).await;
            if tx.send(Ok(initial_update)).await.is_err() {
                debug!("Client disconnected before receiving initial state");
                return;
            }

            // Track last sent progress to enable delta-only streaming
            let mut last_progress = initial_job.progress_percentage;
            let mut last_status = initial_job.status.clone();

            // Stream real-time updates (hybrid approach - step 2)
            loop {
                match status_receiver.recv().await {
                    Ok(update) => {
                        // Delta filtering: only send if meaningful change occurred
                        let progress_changed = (update.progress_percentage - last_progress).abs() > 0.001;
                        let status_changed = update.status != last_status;

                        if progress_changed || status_changed {
                            last_progress = update.progress_percentage;
                            last_status = update.status.clone();

                            let progress_update = Self::convert_status_to_progress(update);

                            // Send update to client
                            if tx.send(Ok(progress_update.clone())).await.is_err() {
                                info!("Client disconnected, closing stream for job {}", job_id);
                                break;
                            }

                            // Close stream if job completed/failed/stopped
                            if Self::is_terminal_status(&progress_update.status) {
                                info!("Job {} reached terminal status, closing stream", job_id);
                                break;
                            }
                        }
                    }
                    Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
                        warn!("Progress stream lagged by {} messages for job {}, re-fetching state", 
                            n, job_id);
                        
                        // Handle lag gracefully: re-fetch current state from orchestrator
                        if let Ok(current_job) = orchestrator.get_job(job_id).await {
                            let recovery_update = Self::create_initial_progress_update(&current_job).await;
                            if tx.send(Ok(recovery_update)).await.is_err() {
                                break;
                            }
                            last_progress = current_job.progress_percentage;
                            last_status = current_job.status;
                        }
                    }
                    Err(_) => {
                        debug!("Broadcast channel closed for job {}", job_id);
                        break;
                    }
                }
            }

            debug!("Progress streaming task terminated for job {}", job_id);
        });

        let stream = ReceiverStream::new(rx);
        Ok(Response::new(Box::pin(stream) as Pin<Box<dyn Stream<Item = Result<ProgressUpdate, Status>> + Send>>))
    }

    /// Create initial progress update from job state
    async fn create_initial_progress_update(job: &crate::orchestrator::TrainingJob) -> ProgressUpdate {
        let status = Self::convert_job_status(&job.status);
        
        ProgressUpdate {
            job_id: job.id.to_string(),
            current_trial: job.current_epoch,
            total_trials: job.total_epochs,
            trial_params: job.tags.clone(),
            trial_sharpe: job.metrics.get("sharpe_ratio").copied().unwrap_or(0.0) as f32,
            best_sharpe_so_far: job.metrics.get("best_sharpe").copied().unwrap_or(0.0) as f32,
            estimated_time_remaining: 0, // Would calculate based on epoch timing
            status: status as i32,
            message: format!("Initial state: {} ({:.1}%)", 
                Self::status_to_string(&job.status), 
                job.progress_percentage),
            timestamp: chrono::Utc::now().timestamp(),
            update_type: UpdateType::UpdateHeartbeat as i32,
        }
    }

    /// Convert internal status update to proto ProgressUpdate
    fn convert_status_to_progress(update: TrainingStatusUpdate) -> ProgressUpdate {
        let status = Self::convert_job_status(&update.status);
        
        // Determine update type based on progress
        let update_type = if update.progress_percentage >= 99.99 {
            UpdateType::UpdateJobComplete
        } else if update.current_epoch > 0 && update.current_epoch < update.total_epochs {
            UpdateType::UpdateTrialComplete
        } else {
            UpdateType::UpdateHeartbeat
        };

        ProgressUpdate {
            job_id: update.job_id.to_string(),
            current_trial: update.current_epoch,
            total_trials: update.total_epochs,
            trial_params: HashMap::new(),
            trial_sharpe: update.metrics.get("sharpe_ratio").copied().unwrap_or(0.0) as f32,
            best_sharpe_so_far: update.metrics.get("best_sharpe").copied().unwrap_or(0.0) as f32,
            estimated_time_remaining: 0,
            status: status as i32,
            message: update.message,
            timestamp: update.timestamp.timestamp(),
            update_type: update_type as i32,
        }
    }

    /// Convert internal job status to proto tuning status
    fn convert_job_status(status: &JobStatus) -> TuningJobStatus {
        match status {
            JobStatus::Pending => TuningJobStatus::TuningPending,
            JobStatus::Running => TuningJobStatus::TuningRunning,
            JobStatus::Completed => TuningJobStatus::TuningCompleted,
            JobStatus::Failed => TuningJobStatus::TuningFailed,
            JobStatus::Stopped => TuningJobStatus::TuningStopped,
            JobStatus::Paused => TuningJobStatus::TuningRunning, // Map paused to running
        }
    }

    /// Check if status is terminal (should close stream)
    fn is_terminal_status(status: &i32) -> bool {
        matches!(
            TuningJobStatus::try_from(*status),
            Ok(TuningJobStatus::TuningCompleted)
                | Ok(TuningJobStatus::TuningFailed)
                | Ok(TuningJobStatus::TuningStopped)
        )
    }

    /// Convert job status to human-readable string
    fn status_to_string(status: &JobStatus) -> &str {
        match status {
            JobStatus::Pending => "Pending",
            JobStatus::Running => "Running",
            JobStatus::Completed => "Completed",
            JobStatus::Failed => "Failed",
            JobStatus::Stopped => "Stopped",
            JobStatus::Paused => "Paused",
        }
    }
}

/// Batch progress streamer for multiplexed jobs
pub struct BatchProgressStreamer {
    orchestrator: Arc<TrainingOrchestrator>,
}

impl BatchProgressStreamer {
    /// Create a new batch progress streamer
    pub fn new(orchestrator: Arc<TrainingOrchestrator>) -> Self {
        Self { orchestrator }
    }

    /// Stream aggregated progress for multiple jobs in a batch
    ///
    /// Supports up to 16 concurrent child jobs with weighted progress calculation.
    /// Multiplexes all job updates into a single stream with overall batch progress.
    pub async fn stream_batch_progress(
        &self,
        batch_id: String,
        job_ids: Vec<Uuid>,
    ) -> Result<Response<Pin<Box<dyn Stream<Item = Result<ProgressUpdate, Status>> + Send>>>, Status>
    {
        if job_ids.is_empty() {
            return Err(Status::invalid_argument("Batch has no jobs"));
        }

        if job_ids.len() > 16 {
            return Err(Status::invalid_argument(
                "Batch size exceeds maximum of 16 jobs",
            ));
        }

        info!(
            "Starting batch progress stream for {} jobs in batch {}",
            job_ids.len(),
            batch_id
        );

        let (tx, rx) = mpsc::channel(100);
        let orchestrator = self.orchestrator.clone();

        tokio::spawn(async move {
            // Track progress for each job
            let mut progress_map: HashMap<Uuid, f32> = HashMap::new();
            let mut status_map: HashMap<Uuid, JobStatus> = HashMap::new();

            // Initialize all jobs to 0% progress
            for job_id in &job_ids {
                progress_map.insert(*job_id, 0.0);
                status_map.insert(*job_id, JobStatus::Pending);
            }

            // Create channel to multiplex all job updates
            let (update_tx, mut update_rx) = mpsc::channel(100);

            // Spawn receiver task for each job
            for job_id in job_ids.clone() {
                let update_tx_clone = update_tx.clone();
                let orchestrator_clone = orchestrator.clone();

                tokio::spawn(async move {
                    // Fetch initial state
                    if let Ok(job) = orchestrator_clone.get_job(job_id).await {
                        let initial_update = TrainingStatusUpdate {
                            job_id: job.id,
                            status: job.status.clone(),
                            progress_percentage: job.progress_percentage,
                            current_epoch: job.current_epoch,
                            total_epochs: job.total_epochs,
                            metrics: HashMap::new(),
                            message: "Initial state".to_string(),
                            timestamp: chrono::Utc::now(),
                            financial_metrics: None,
                            resource_usage: Default::default(),
                        };
                        let _ = update_tx_clone.send(initial_update).await;
                    }

                    // Subscribe to live updates
                    if let Ok(mut receiver) = orchestrator_clone.subscribe_to_job_status(job_id).await {
                        while let Ok(update) = receiver.recv().await {
                            if update_tx_clone.send(update).await.is_err() {
                                break;
                            }
                        }
                    }
                });
            }
            drop(update_tx);

            // Aggregate updates from all jobs
            while let Some(update) = update_rx.recv().await {
                // Update individual job progress
                progress_map.insert(update.job_id, update.progress_percentage);
                status_map.insert(update.job_id, update.status.clone());

                // Calculate weighted overall batch progress (equal weight)
                let overall_progress: f32 =
                    progress_map.values().sum::<f32>() / job_ids.len() as f32;

                // Create batch progress update
                let batch_update = ProgressUpdate {
                    job_id: batch_id.clone(),
                    current_trial: progress_map.len() as u32,
                    total_trials: job_ids.len() as u32,
                    trial_params: HashMap::new(),
                    trial_sharpe: 0.0,
                    best_sharpe_so_far: 0.0,
                    estimated_time_remaining: 0,
                    status: TuningJobStatus::TuningRunning as i32,
                    message: format!(
                        "Batch progress: {:.1}% ({}/{} jobs)",
                        overall_progress,
                        progress_map.values().filter(|&&p| p >= 100.0).count(),
                        job_ids.len()
                    ),
                    timestamp: chrono::Utc::now().timestamp(),
                    update_type: UpdateType::UpdateHeartbeat as i32,
                };

                if tx.send(Ok(batch_update)).await.is_err() {
                    info!("Client disconnected from batch stream {}", batch_id);
                    break;
                }

                // Check if all jobs completed
                let all_complete = status_map
                    .values()
                    .all(|s| matches!(s, JobStatus::Completed | JobStatus::Failed | JobStatus::Stopped));

                if all_complete {
                    info!("All jobs in batch {} completed", batch_id);
                    break;
                }
            }

            debug!("Batch progress streaming task terminated for batch {}", batch_id);
        });

        let stream = ReceiverStream::new(rx);
        Ok(Response::new(Box::pin(stream) as Pin<Box<dyn Stream<Item = Result<ProgressUpdate, Status>> + Send>>))
    }
}

impl Default for crate::orchestrator::ResourceUsage {
    fn default() -> Self {
        Self {
            cpu_usage_percent: 0.0,
            memory_usage_gb: 0.0,
            gpu_usage_percent: None,
            gpu_memory_usage_gb: None,
            active_workers: 0,
        }
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;

    #[test]
    fn test_status_conversion() {
        assert_eq!(
            ProgressStreamer::convert_job_status(&JobStatus::Pending),
            TuningJobStatus::TuningPending
        );
        assert_eq!(
            ProgressStreamer::convert_job_status(&JobStatus::Running),
            TuningJobStatus::TuningRunning
        );
        assert_eq!(
            ProgressStreamer::convert_job_status(&JobStatus::Completed),
            TuningJobStatus::TuningCompleted
        );
        assert_eq!(
            ProgressStreamer::convert_job_status(&JobStatus::Failed),
            TuningJobStatus::TuningFailed
        );
        assert_eq!(
            ProgressStreamer::convert_job_status(&JobStatus::Stopped),
            TuningJobStatus::TuningStopped
        );
    }

    #[test]
    fn test_is_terminal_status() {
        assert!(ProgressStreamer::is_terminal_status(&(TuningJobStatus::TuningCompleted as i32)));
        assert!(ProgressStreamer::is_terminal_status(&(TuningJobStatus::TuningFailed as i32)));
        assert!(ProgressStreamer::is_terminal_status(&(TuningJobStatus::TuningStopped as i32)));
        assert!(!ProgressStreamer::is_terminal_status(&(TuningJobStatus::TuningRunning as i32)));
        assert!(!ProgressStreamer::is_terminal_status(&(TuningJobStatus::TuningPending as i32)));
    }

    #[test]
    fn test_status_to_string() {
        assert_eq!(ProgressStreamer::status_to_string(&JobStatus::Pending), "Pending");
        assert_eq!(ProgressStreamer::status_to_string(&JobStatus::Running), "Running");
        assert_eq!(ProgressStreamer::status_to_string(&JobStatus::Completed), "Completed");
        assert_eq!(ProgressStreamer::status_to_string(&JobStatus::Failed), "Failed");
        assert_eq!(ProgressStreamer::status_to_string(&JobStatus::Stopped), "Stopped");
        assert_eq!(ProgressStreamer::status_to_string(&JobStatus::Paused), "Paused");
    }
}
