//! gRPC Hyperparameter Tuning Handlers
//!
//! This module implements the gRPC handlers for hyperparameter tuning operations,
//! integrating with the TuningManager to coordinate Optuna subprocess execution.

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

use async_stream::stream;
use tokio_stream::Stream;
use tonic::{Request, Response, Status};
use tracing::{debug, info, warn};
use uuid::Uuid;

use crate::service::proto::{
    GetTuningJobStatusRequest, GetTuningJobStatusResponse, ProgressUpdate, StartTuningJobRequest,
    StartTuningJobResponse, StopTuningJobRequest, StopTuningJobResponse, StreamProgressRequest,
    TrialResult as ProtoTrialResult, TrialState as ProtoTrialState,
    TuningJobStatus as ProtoTuningJobStatus, UpdateType,
};
use crate::tuning_manager::{ProgressUpdateType, TrialState, TuningJobStatus, TuningManager};

/// Implementation of hyperparameter tuning gRPC handlers
pub struct TuningHandlers {
    tuning_manager: Arc<TuningManager>,
}

impl TuningHandlers {
    /// Create new tuning handlers
    pub fn new(tuning_manager: Arc<TuningManager>) -> Self {
        Self { tuning_manager }
    }

    /// Start a new hyperparameter tuning job
    pub async fn start_tuning_job(
        &self,
        request: Request<StartTuningJobRequest>,
    ) -> Result<Response<StartTuningJobResponse>, Status> {
        let req = request.into_inner();

        info!(
            "Starting hyperparameter tuning job for model type: {} with {} trials",
            req.model_type, req.num_trials
        );

        // Validate input
        if req.model_type.is_empty() {
            return Err(Status::invalid_argument("Model type cannot be empty"));
        }

        if req.num_trials == 0 {
            return Err(Status::invalid_argument("Number of trials must be > 0"));
        }

        if req.config_path.is_empty() {
            return Err(Status::invalid_argument(
                "Configuration path cannot be empty",
            ));
        }

        // Extract tags
        let tags: HashMap<String, String> = req.tags.into_iter().collect();

        // Submit tuning job to manager
        let job_id = self
            .tuning_manager
            .start_tuning_job(
                req.model_type,
                req.num_trials,
                req.config_path,
                req.description,
                tags,
            )
            .await
            .map_err(|e| Status::internal(format!("Failed to start tuning job: {}", e)))?;

        info!(
            "Hyperparameter tuning job submitted successfully: {}",
            job_id
        );

        let response = StartTuningJobResponse {
            job_id: job_id.to_string(),
            status: ProtoTuningJobStatus::TuningRunning as i32,
            message: "Hyperparameter tuning job started successfully".to_string(),
        };

        Ok(Response::new(response))
    }

    /// Get status of a hyperparameter tuning job
    pub async fn get_tuning_job_status(
        &self,
        request: Request<GetTuningJobStatusRequest>,
    ) -> Result<Response<GetTuningJobStatusResponse>, Status> {
        let req = request.into_inner();

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

        debug!("Getting status for tuning job: {}", job_id);

        // Get job status from manager
        let job = self
            .tuning_manager
            .get_tuning_job_status(job_id)
            .await
            .map_err(|e| Status::not_found(format!("Tuning job not found: {}", e)))?;

        // Convert to protobuf format
        let proto_status = Self::convert_tuning_status(&job.status);

        let trial_history: Vec<ProtoTrialResult> = job
            .trial_history
            .iter()
            .map(|trial| ProtoTrialResult {
                trial_number: trial.trial_number,
                params: trial.params.clone(),
                objective_value: trial.objective_value,
                metrics: trial.metrics.clone(),
                state: Self::convert_trial_state(&trial.state) as i32,
                started_at: trial.started_at.timestamp(),
                completed_at: trial.completed_at.map(|dt| dt.timestamp()).unwrap_or(0),
            })
            .collect();

        let response = GetTuningJobStatusResponse {
            job_id: job.id.to_string(),
            status: proto_status as i32,
            current_trial: job.current_trial,
            total_trials: job.num_trials,
            best_params: job.best_params,
            best_metrics: job.best_metrics,
            trial_history,
            message: job.error_message.unwrap_or_else(|| "Running".to_string()),
            started_at: job.started_at.timestamp(),
            updated_at: job.updated_at.timestamp(),
        };

        Ok(Response::new(response))
    }

    /// Stop a running hyperparameter tuning job
    pub async fn stop_tuning_job(
        &self,
        request: Request<StopTuningJobRequest>,
    ) -> Result<Response<StopTuningJobResponse>, Status> {
        let req = request.into_inner();

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

        info!("Stopping tuning job: {} (reason: {})", job_id, req.reason);

        // Stop the tuning job
        self.tuning_manager
            .stop_tuning_job(job_id, req.reason.clone())
            .await
            .map_err(|e| Status::internal(format!("Failed to stop tuning job: {}", e)))?;

        // Get final status
        let job = self
            .tuning_manager
            .get_tuning_job_status(job_id)
            .await
            .map_err(|e| Status::not_found(format!("Tuning job not found: {}", e)))?;

        let final_status = Self::convert_tuning_status(&job.status);

        let response = StopTuningJobResponse {
            success: true,
            message: format!(
                "Tuning job stopped successfully. Completed {} of {} trials.",
                job.current_trial, job.num_trials
            ),
            final_status: final_status as i32,
        };

        Ok(Response::new(response))
    }

    /// Convert internal tuning status to protobuf status
    fn convert_tuning_status(status: &TuningJobStatus) -> ProtoTuningJobStatus {
        match status {
            TuningJobStatus::Pending => ProtoTuningJobStatus::TuningPending,
            TuningJobStatus::Running => ProtoTuningJobStatus::TuningRunning,
            TuningJobStatus::Completed => ProtoTuningJobStatus::TuningCompleted,
            TuningJobStatus::Failed => ProtoTuningJobStatus::TuningFailed,
            TuningJobStatus::Stopped => ProtoTuningJobStatus::TuningStopped,
        }
    }

    /// Convert internal trial state to protobuf state
    fn convert_trial_state(state: &TrialState) -> ProtoTrialState {
        match state {
            TrialState::Running => ProtoTrialState::TrialRunning,
            TrialState::Complete => ProtoTrialState::TrialComplete,
            TrialState::Pruned => ProtoTrialState::TrialPruned,
            TrialState::Failed => ProtoTrialState::TrialFailed,
        }
    }

    /// Stream real-time tuning progress updates
    pub async fn stream_tuning_progress(
        &self,
        request: Request<StreamProgressRequest>,
    ) -> Result<Response<Pin<Box<dyn Stream<Item = Result<ProgressUpdate, Status>> + Send>>>, Status>
    {
        let req = request.into_inner();

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

        info!("Client subscribing to progress stream for job: {}", job_id);

        // Verify job exists
        self.tuning_manager
            .get_tuning_job_status(job_id)
            .await
            .map_err(|e| Status::not_found(format!("Tuning job not found: {}", e)))?;

        // Subscribe to progress updates
        let mut rx = self.tuning_manager.subscribe_to_progress(job_id);

        // Create stream
        let output_stream = stream! {
            loop {
                match rx.recv().await {
                    Ok(event) => {
                        // Filter for this specific job_id
                        if event.job_id != job_id {
                            continue;
                        }

                        // Convert to protobuf message
                        let update_type = match event.update_type {
                            ProgressUpdateType::TrialComplete => UpdateType::UpdateTrialComplete,
                            ProgressUpdateType::Heartbeat => UpdateType::UpdateHeartbeat,
                            ProgressUpdateType::JobComplete => UpdateType::UpdateJobComplete,
                        };

                        // Convert trial params to strings
                        let trial_params: HashMap<String, String> = event.trial_params
                            .into_iter()
                            .map(|(k, v)| (k, format!("{:.6}", v)))
                            .collect();

                        let progress_update = ProgressUpdate {
                            job_id: event.job_id.to_string(),
                            current_trial: event.current_trial,
                            total_trials: event.total_trials,
                            trial_params,
                            trial_sharpe: event.trial_sharpe,
                            best_sharpe_so_far: event.best_sharpe_so_far,
                            estimated_time_remaining: event.estimated_time_remaining,
                            status: Self::convert_tuning_status(&event.status) as i32,
                            message: event.message,
                            timestamp: event.timestamp.timestamp(),
                            update_type: update_type as i32,
                        };

                        yield Ok(progress_update);

                        // Close stream if job is complete
                        if event.update_type == ProgressUpdateType::JobComplete {
                            info!("Job {} complete, closing progress stream", job_id);
                            break;
                        }
                    }
                    Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => {
                        warn!("Progress stream lagged, skipped {} messages", skipped);
                        // Continue streaming despite lag
                    }
                    Err(tokio::sync::broadcast::error::RecvError::Closed) => {
                        info!("Progress channel closed for job {}", job_id);
                        break;
                    }
                }
            }
        };

        Ok(Response::new(Box::pin(output_stream)))
    }
}

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

    #[test]
    fn test_status_conversion() {
        assert_eq!(
            TuningHandlers::convert_tuning_status(&TuningJobStatus::Pending),
            ProtoTuningJobStatus::TuningPending
        );

        assert_eq!(
            TuningHandlers::convert_tuning_status(&TuningJobStatus::Running),
            ProtoTuningJobStatus::TuningRunning
        );

        assert_eq!(
            TuningHandlers::convert_tuning_status(&TuningJobStatus::Completed),
            ProtoTuningJobStatus::TuningCompleted
        );

        assert_eq!(
            TuningHandlers::convert_tuning_status(&TuningJobStatus::Failed),
            ProtoTuningJobStatus::TuningFailed
        );

        assert_eq!(
            TuningHandlers::convert_tuning_status(&TuningJobStatus::Stopped),
            ProtoTuningJobStatus::TuningStopped
        );
    }

    #[test]
    fn test_trial_state_conversion() {
        assert_eq!(
            TuningHandlers::convert_trial_state(&TrialState::Running),
            ProtoTrialState::TrialRunning
        );

        assert_eq!(
            TuningHandlers::convert_trial_state(&TrialState::Complete),
            ProtoTrialState::TrialComplete
        );

        assert_eq!(
            TuningHandlers::convert_trial_state(&TrialState::Pruned),
            ProtoTrialState::TrialPruned
        );

        assert_eq!(
            TuningHandlers::convert_trial_state(&TrialState::Failed),
            ProtoTrialState::TrialFailed
        );
    }

    #[test]
    fn test_tuning_job_validation() {
        let mut tags = HashMap::new();
        tags.insert("env".to_string(), "test".to_string());

        let job = TuningJob::new("TLOB".to_string(), 100, "Test tuning job".to_string(), tags);

        assert_eq!(job.model_type, "TLOB");
        assert_eq!(job.num_trials, 100);
        assert_eq!(job.current_trial, 0);
        assert!(job.best_params.is_empty());
        assert!(job.trial_history.is_empty());
    }
}
