//! Hyperparameter Tuning Manager
//!
//! This module manages hyperparameter tuning jobs using Optuna subprocess.
//! It spawns Python Optuna processes, tracks their status, and coordinates
//! trial execution with the ML training service.

use std::collections::HashMap;
use std::process::{Child, Command, Stdio};
use std::sync::Arc;
use std::time::Duration;

use anyhow::{anyhow, Context, Result};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use tokio::fs;
use tokio::sync::{broadcast, RwLock};
use tracing::{debug, error, info, warn};
use uuid::Uuid;

/// Status of a hyperparameter tuning job
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum TuningJobStatus {
    Pending,
    Running,
    Completed,
    Failed,
    Stopped,
}

/// Individual trial result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrialResult {
    pub trial_number: u32,
    pub params: HashMap<String, f32>,
    pub objective_value: f32,
    pub metrics: HashMap<String, f32>,
    pub state: TrialState,
    pub started_at: DateTime<Utc>,
    pub completed_at: Option<DateTime<Utc>>,
}

/// Trial execution state
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum TrialState {
    Running,
    Complete,
    Pruned,
    Failed,
}

/// Tuning job metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TuningJob {
    pub id: Uuid,
    pub model_type: String,
    pub status: TuningJobStatus,
    pub num_trials: u32,
    pub current_trial: u32,
    pub best_params: HashMap<String, f32>,
    pub best_metrics: HashMap<String, f32>,
    pub trial_history: Vec<TrialResult>,
    pub started_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
    pub completed_at: Option<DateTime<Utc>>,
    pub description: String,
    pub tags: HashMap<String, String>,
    pub error_message: Option<String>,
}

impl TuningJob {
    pub fn new(
        model_type: String,
        num_trials: u32,
        description: String,
        tags: HashMap<String, String>,
    ) -> Self {
        let now = Utc::now();
        Self {
            id: Uuid::new_v4(),
            model_type,
            status: TuningJobStatus::Pending,
            num_trials,
            current_trial: 0,
            best_params: HashMap::new(),
            best_metrics: HashMap::new(),
            trial_history: Vec::new(),
            started_at: now,
            updated_at: now,
            completed_at: None,
            description,
            tags,
            error_message: None,
        }
    }
}

#[async_trait]
impl TuningManagerTrait for TuningManager {
    async fn start_tuning_job(
        &self,
        model_type: String,
        num_trials: u32,
        config_path: String,
        description: String,
        tags: HashMap<String, String>,
    ) -> Result<Uuid> {
        self.start_tuning_job(model_type, num_trials, config_path, description, tags)
            .await
    }

    async fn get_tuning_job_status(&self, job_id: Uuid) -> Result<TuningJob> {
        self.get_tuning_job_status(job_id).await
    }

    async fn stop_tuning_job(&self, job_id: Uuid, reason: String) -> Result<()> {
        self.stop_tuning_job(job_id, reason).await
    }
}

/// Process handle for running Optuna subprocess
struct ProcessHandle {
    child: Child,
}

/// Progress update event for streaming
#[derive(Debug, Clone)]
pub struct ProgressUpdateEvent {
    pub job_id: Uuid,
    pub current_trial: u32,
    pub total_trials: u32,
    pub trial_params: HashMap<String, f32>,
    pub trial_sharpe: f32,
    pub best_sharpe_so_far: f32,
    pub estimated_time_remaining: u32,
    pub status: TuningJobStatus,
    pub message: String,
    pub timestamp: DateTime<Utc>,
    pub update_type: ProgressUpdateType,
}

/// Type of progress update
#[derive(Debug, Clone, PartialEq)]
pub enum ProgressUpdateType {
    TrialComplete,
    Heartbeat,
    JobComplete,
}

/// Trait abstraction for TuningManager to enable dependency injection and testing
#[async_trait]
pub trait TuningManagerTrait: Send + Sync {
    /// Start a new hyperparameter tuning job
    async fn start_tuning_job(
        &self,
        model_type: String,
        num_trials: u32,
        config_path: String,
        description: String,
        tags: HashMap<String, String>,
    ) -> Result<Uuid>;

    /// Get status of a tuning job
    async fn get_tuning_job_status(&self, job_id: Uuid) -> Result<TuningJob>;

    /// Stop a running tuning job
    async fn stop_tuning_job(&self, job_id: Uuid, reason: String) -> Result<()>;
}

/// Tuning manager coordinates hyperparameter tuning jobs
pub struct TuningManager {
    /// Active tuning jobs indexed by job_id
    jobs: Arc<RwLock<HashMap<Uuid, TuningJob>>>,

    /// Running subprocess handles
    processes: Arc<RwLock<HashMap<Uuid, ProcessHandle>>>,

    /// Path to Python hyperparameter tuner script
    tuner_script_path: String,

    /// Working directory for tuning artifacts
    working_dir: String,

    /// Broadcast channel for progress updates (capacity: 100 messages)
    progress_tx: broadcast::Sender<ProgressUpdateEvent>,
}

impl TuningManager {
    /// Create a new tuning manager
    pub fn new(tuner_script_path: String, working_dir: String) -> Self {
        let (progress_tx, _) = broadcast::channel(100);
        Self {
            jobs: Arc::new(RwLock::new(HashMap::new())),
            processes: Arc::new(RwLock::new(HashMap::new())),
            tuner_script_path,
            working_dir,
            progress_tx,
        }
    }

    /// Subscribe to progress updates for a specific job
    pub fn subscribe_to_progress(&self, _job_id: Uuid) -> broadcast::Receiver<ProgressUpdateEvent> {
        

        // Return a receiver that filters for this job_id
        // Note: Filtering happens in the stream handler
        self.progress_tx.subscribe()
    }

    /// Publish a progress update event
    fn _publish_progress(&self, event: ProgressUpdateEvent) {
        // Best-effort send - warn if no subscribers
        if let Err(e) = self.progress_tx.send(event) {
            tracing::warn!("Channel send failed (receiver dropped?): {:?}", e);
        }
    }

    /// Start a new hyperparameter tuning job
    pub async fn start_tuning_job(
        &self,
        model_type: String,
        num_trials: u32,
        config_path: String,
        description: String,
        tags: HashMap<String, String>,
    ) -> Result<Uuid> {
        // Create job metadata
        let mut job = TuningJob::new(model_type.clone(), num_trials, description, tags);
        let job_id = job.id;

        info!(
            "Starting hyperparameter tuning job {} for model {} with {} trials",
            job_id, model_type, num_trials
        );

        // Update job status to running
        job.status = TuningJobStatus::Running;
        job.updated_at = Utc::now();

        // Store job metadata
        {
            let mut jobs = self.jobs.write().await;
            jobs.insert(job_id, job.clone());
        }

        // Spawn Optuna subprocess asynchronously
        let tuner_script = self.tuner_script_path.clone();
        let working_dir = self.working_dir.clone();
        let processes = Arc::clone(&self.processes);
        let jobs = Arc::clone(&self.jobs);
        let progress_tx = self.progress_tx.clone();

        tokio::spawn(async move {
            match Self::spawn_optuna_process(
                job_id,
                model_type,
                num_trials,
                config_path,
                tuner_script,
                working_dir.clone(),
            )
            .await
            {
                Ok(child) => {
                    // Store process handle
                    let mut procs = processes.write().await;
                    procs.insert(job_id, ProcessHandle { child });
                    info!("Optuna subprocess spawned for job {}", job_id);

                    // Monitor process completion
                    drop(procs); // Release lock before monitoring
                    Self::monitor_process(
                        job_id,
                        processes.clone(),
                        jobs.clone(),
                        working_dir,
                        progress_tx,
                    )
                    .await;
                },
                Err(e) => {
                    error!("Failed to spawn Optuna process for job {}: {}", job_id, e);
                    let mut jobs_guard = jobs.write().await;
                    if let Some(job) = jobs_guard.get_mut(&job_id) {
                        job.status = TuningJobStatus::Failed;
                        job.error_message = Some(format!("Failed to spawn process: {}", e));
                        job.updated_at = Utc::now();
                    }
                },
            }
        });

        Ok(job_id)
    }

    /// Get status of a tuning job
    pub async fn get_tuning_job_status(&self, job_id: Uuid) -> Result<TuningJob> {
        // Try to load latest status from disk
        if let Ok(job) = self.load_job_status(job_id).await {
            // Update in-memory state
            let mut jobs = self.jobs.write().await;
            jobs.insert(job_id, job.clone());
            return Ok(job);
        }

        // Fall back to in-memory state
        let jobs = self.jobs.read().await;
        jobs.get(&job_id)
            .cloned()
            .ok_or_else(|| anyhow!("Tuning job {} not found", job_id))
    }

    /// Stop a running tuning job
    pub async fn stop_tuning_job(&self, job_id: Uuid, reason: String) -> Result<()> {
        info!("Stopping tuning job {}: {}", job_id, reason);

        // Get process handle
        let mut processes = self.processes.write().await;
        if let Some(mut handle) = processes.remove(&job_id) {
            // Send SIGTERM for graceful shutdown
            #[cfg(unix)]
            {
                use nix::sys::signal::{kill, Signal};
                use nix::unistd::Pid;

                let pid = Pid::from_raw(handle.child.id() as i32);
                if let Err(e) = kill(pid, Signal::SIGTERM) {
                    warn!("Failed to send SIGTERM to process {}: {}", pid, e);
                }
            }

            #[cfg(not(unix))]
            {
                if let Err(e) = handle.child.kill() {
                    warn!("Failed to kill process: {}", e);
                }
            }

            // Wait for graceful shutdown with timeout
            let timeout_duration = Duration::from_secs(5);
            let shutdown_result = tokio::time::timeout(
                timeout_duration,
                tokio::task::spawn_blocking(move || handle.child.wait()),
            )
            .await;

            match shutdown_result {
                Ok(Ok(Ok(status))) => {
                    info!("Process exited with status: {}", status);
                },
                Ok(Ok(Err(e))) => {
                    warn!("Error waiting for process: {}", e);
                },
                Ok(Err(e)) => {
                    warn!("Join error waiting for process: {}", e);
                },
                Err(_) => {
                    warn!(
                        "Process did not exit within {}s, may still be running",
                        timeout_duration.as_secs()
                    );
                },
            }
        } else {
            debug!("No running process found for job {}", job_id);
        }

        // Update job status
        let mut jobs = self.jobs.write().await;
        if let Some(job) = jobs.get_mut(&job_id) {
            job.status = TuningJobStatus::Stopped;
            job.updated_at = Utc::now();
            job.completed_at = Some(Utc::now());
        }

        Ok(())
    }

    /// Spawn Optuna hyperparameter tuning process
    async fn spawn_optuna_process(
        job_id: Uuid,
        model_type: String,
        num_trials: u32,
        config_path: String,
        tuner_script: String,
        working_dir: String,
    ) -> Result<Child> {
        // Create job-specific output directory
        let job_dir = format!("{}/{}", working_dir, job_id);
        fs::create_dir_all(&job_dir)
            .await
            .context("Failed to create job directory")?;

        // Spawn Python Optuna subprocess
        let child = Command::new("python3")
            .arg(&tuner_script)
            .arg("--job-id")
            .arg(job_id.to_string())
            .arg("--model-type")
            .arg(model_type)
            .arg("--num-trials")
            .arg(num_trials.to_string())
            .arg("--config-path")
            .arg(config_path)
            .arg("--output-dir")
            .arg(&job_dir)
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .context("Failed to spawn Optuna process")?;

        Ok(child)
    }

    /// Monitor process completion and update job status
    async fn monitor_process(
        job_id: Uuid,
        processes: Arc<RwLock<HashMap<Uuid, ProcessHandle>>>,
        jobs: Arc<RwLock<HashMap<Uuid, TuningJob>>>,
        working_dir: String,
        progress_tx: broadcast::Sender<ProgressUpdateEvent>,
    ) {
        let mut last_trial = 0u32;
        let start_time = Utc::now();
        let mut heartbeat_counter = 0u32;

        // Poll for process completion
        loop {
            tokio::time::sleep(Duration::from_secs(5)).await;
            heartbeat_counter += 1;

            // Check if process still exists
            let still_running = {
                let procs = processes.read().await;
                procs.contains_key(&job_id)
            };

            if !still_running {
                break;
            }

            // Try to load updated status from disk
            let status_path = format!("{}/{}/status.json", working_dir, job_id);
            if let Ok(contents) = fs::read_to_string(&status_path).await {
                if let Ok(job_update) = serde_json::from_str::<TuningJob>(&contents) {
                    let mut jobs_guard = jobs.write().await;
                    if let Some(job) = jobs_guard.get_mut(&job_id) {
                        let trial_changed = job_update.current_trial > last_trial;

                        job.current_trial = job_update.current_trial;
                        job.best_params = job_update.best_params.clone();
                        job.best_metrics = job_update.best_metrics.clone();
                        job.trial_history = job_update.trial_history.clone();
                        job.updated_at = Utc::now();

                        // Calculate estimated time remaining
                        let elapsed =
                            Utc::now().signed_duration_since(start_time).num_seconds() as u32;
                        let avg_trial_time = if job.current_trial > 0 {
                            elapsed / job.current_trial
                        } else {
                            0
                        };
                        let remaining_trials = job.num_trials.saturating_sub(job.current_trial);
                        let estimated_time_remaining = avg_trial_time * remaining_trials;

                        // Get best Sharpe ratio
                        let best_sharpe =
                            job.best_metrics.get("sharpe_ratio").copied().unwrap_or(0.0);

                        // Get current trial Sharpe ratio
                        let trial_sharpe = job
                            .trial_history
                            .last()
                            .map(|t| t.objective_value)
                            .unwrap_or(0.0);

                        // Get current trial params
                        let trial_params = job
                            .trial_history
                            .last()
                            .map(|t| t.params.clone())
                            .unwrap_or_default();

                        // Determine update type
                        let update_type = if trial_changed {
                            last_trial = job.current_trial;
                            ProgressUpdateType::TrialComplete
                        } else if heartbeat_counter % 6 == 0 {
                            // Heartbeat every 30s
                            ProgressUpdateType::Heartbeat
                        } else {
                            continue; // Skip this iteration
                        };

                        // Publish progress update
                        let event = ProgressUpdateEvent {
                            job_id,
                            current_trial: job.current_trial,
                            total_trials: job.num_trials,
                            trial_params,
                            trial_sharpe,
                            best_sharpe_so_far: best_sharpe,
                            estimated_time_remaining,
                            status: job.status.clone(),
                            message: format!(
                                "Trial {}/{} complete",
                                job.current_trial, job.num_trials
                            ),
                            timestamp: Utc::now(),
                            update_type,
                        };

                        let _ = progress_tx.send(event);

                        // Check if completed
                        if job_update.status == TuningJobStatus::Completed
                            || job_update.status == TuningJobStatus::Failed
                        {
                            job.status = job_update.status;
                            job.completed_at = Some(Utc::now());
                            job.error_message = job_update.error_message;

                            // Remove from processes
                            let mut procs = processes.write().await;
                            procs.remove(&job_id);
                            break;
                        }
                    }
                }
            }
        }

        info!("Monitoring completed for tuning job {}", job_id);
    }

    /// Load job status from disk (written by Optuna subprocess)
    async fn load_job_status(&self, job_id: Uuid) -> Result<TuningJob> {
        let status_path = format!("{}/{}/status.json", self.working_dir, job_id);
        let contents = fs::read_to_string(&status_path)
            .await
            .context("Failed to read status file")?;

        let job: TuningJob =
            serde_json::from_str(&contents).context("Failed to parse status file")?;

        Ok(job)
    }
}

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

    #[test]
    fn test_tuning_job_creation() {
        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.clone(),
        );

        assert_eq!(job.model_type, "TLOB");
        assert_eq!(job.status, TuningJobStatus::Pending);
        assert_eq!(job.num_trials, 100);
        assert_eq!(job.current_trial, 0);
        assert_eq!(job.description, "Test tuning job");
        assert_eq!(job.tags.get("env"), Some(&"test".to_string()));
        assert!(job.best_params.is_empty());
        assert!(job.trial_history.is_empty());
    }

    #[test]
    fn test_trial_result_creation() {
        let mut params = HashMap::new();
        params.insert("learning_rate".to_string(), 0.001);
        params.insert("batch_size".to_string(), 64.0);

        let mut metrics = HashMap::new();
        metrics.insert("sharpe_ratio".to_string(), 1.85);
        metrics.insert("training_loss".to_string(), 0.042);

        let trial = TrialResult {
            trial_number: 1,
            params,
            objective_value: 1.85,
            metrics,
            state: TrialState::Complete,
            started_at: Utc::now(),
            completed_at: Some(Utc::now()),
        };

        assert_eq!(trial.trial_number, 1);
        assert_eq!(trial.objective_value, 1.85);
        assert_eq!(trial.state, TrialState::Complete);
        assert!(trial.completed_at.is_some());
    }

    #[tokio::test]
    async fn test_tuning_manager_creation() {
        let manager = TuningManager::new(
            "/path/to/hyperparameter_tuner.py".to_string(),
            "/tmp/tuning_jobs".to_string(),
        );

        let jobs = manager.jobs.read().await;
        assert_eq!(jobs.len(), 0);
    }

    #[tokio::test]
    async fn test_get_nonexistent_job() {
        let manager = TuningManager::new(
            "/path/to/hyperparameter_tuner.py".to_string(),
            "/tmp/tuning_jobs".to_string(),
        );

        let job_id = Uuid::new_v4();
        let result = manager.get_tuning_job_status(job_id).await;
        assert!(result.is_err());
    }
}
