//! `fxt train` -- ML training lifecycle.

use anyhow::Result;
use clap::{Parser, Subcommand};
use colored::Colorize;
use serde::Serialize;
use std::io::{self, Write};
use tokio_stream::StreamExt;

use crate::grpc::FoxhuntClient;
use crate::output::{self, HumanReadable, OutputFormat};
use crate::proto::{ml_training, monitoring};

#[derive(Parser, Debug)]
pub struct TrainCommand {
    #[command(subcommand)]
    action: TrainAction,
}

#[derive(Subcommand, Debug)]
enum TrainAction {
    /// Start a training job
    Start {
        /// Model type (e.g. dqn, ppo, tft, mamba2)
        #[arg(long)]
        model: String,
        /// Path to training config YAML
        #[arg(long)]
        config: Option<String>,
        /// Enable GPU training
        #[arg(long)]
        gpu: bool,
    },
    /// Stop a running training job
    Stop {
        /// Training job ID
        job_id: String,
    },
    /// Show training job status
    Status {
        /// Training job ID
        job_id: String,
    },
    /// List training jobs
    List {
        /// Filter by status (running, completed, failed)
        #[arg(long)]
        status: Option<String>,
        /// Filter by model type
        #[arg(long)]
        model: Option<String>,
    },
    /// Stream training logs
    Logs {
        /// Training job ID
        job_id: String,
        /// Follow log output
        #[arg(long, short)]
        follow: bool,
    },
    /// Live training metrics from all active sessions (via monitoring service)
    #[clap(
        long_about = "Stream live training metrics from the monitoring service.\n\n\
        Shows per-model epoch, loss, validation loss, throughput, GPU utilization,\n\
        financial metrics (Sharpe, Sortino, win rate, max drawdown), and health\n\
        counters (NaN, gradient explosions, checkpoints).\n\n\
        Examples:\n\
          fxt train monitor\n\
          fxt train monitor --once\n\
          fxt train monitor --model dqn --interval 5"
    )]
    Monitor {
        /// Print a single snapshot and exit (no live streaming)
        #[arg(long)]
        once: bool,
        /// Filter to a specific model (e.g. dqn, ppo, tft)
        #[arg(long)]
        model: Option<String>,
        /// Streaming interval in seconds
        #[arg(long, default_value = "3")]
        interval: u32,
    },
}

// ── Result types ──────────────────────────────────────────────────────

#[derive(Serialize)]
struct TrainStartResult {
    job_id: String,
    status: String,
    message: String,
}

#[derive(Serialize)]
struct TrainStopResult {
    success: bool,
    message: String,
}

#[derive(Serialize)]
struct TrainStatusResult {
    job_id: String,
    model_type: String,
    status: String,
    description: String,
    created_at: i64,
    started_at: i64,
    completed_at: i64,
    error: String,
    artifact_path: String,
    financial_metrics: Option<FinancialMetricsSummary>,
}

#[derive(Serialize)]
struct FinancialMetricsSummary {
    sharpe_ratio: f32,
    max_drawdown: f32,
    hit_rate: f32,
    simulated_return: f32,
}

#[derive(Serialize)]
struct TrainListResult {
    total: u32,
    jobs: Vec<TrainJobRow>,
}

#[derive(Serialize)]
struct TrainJobRow {
    job_id: String,
    model_type: String,
    status: String,
    description: String,
    final_loss: f32,
    best_val: f32,
    created_at: i64,
}

// ── Display helpers ──────────────────────────────────────────────────

fn training_status_str(s: i32) -> &'static str {
    use crate::proto::ml_training::TrainingStatus;
    match TrainingStatus::try_from(s) {
        Ok(TrainingStatus::Pending) => "pending",
        Ok(TrainingStatus::Running) => "running",
        Ok(TrainingStatus::Completed) => "completed",
        Ok(TrainingStatus::Failed) => "failed",
        Ok(TrainingStatus::Stopped) => "stopped",
        Ok(TrainingStatus::Paused) => "paused",
        _ => "unknown",
    }
}

fn colorize_status(s: &str) -> String {
    match s {
        "running" => s.cyan().bold().to_string(),
        "completed" => s.green().bold().to_string(),
        "failed" => s.red().bold().to_string(),
        "stopped" => s.yellow().to_string(),
        "pending" => s.dimmed().to_string(),
        "paused" => s.yellow().dimmed().to_string(),
        _ => s.dimmed().to_string(),
    }
}

fn status_filter_to_i32(s: &str) -> i32 {
    use crate::proto::ml_training::TrainingStatus;
    match s.to_lowercase().as_str() {
        "pending" => TrainingStatus::Pending as i32,
        "running" => TrainingStatus::Running as i32,
        "completed" => TrainingStatus::Completed as i32,
        "failed" => TrainingStatus::Failed as i32,
        "stopped" => TrainingStatus::Stopped as i32,
        "paused" => TrainingStatus::Paused as i32,
        _ => TrainingStatus::Unknown as i32,
    }
}

fn format_timestamp(ts: i64) -> String {
    if ts == 0 {
        return "-".to_owned();
    }
    chrono::DateTime::from_timestamp(ts, 0)
        .map(|dt| dt.format("%Y-%m-%d %H:%M:%S").to_string())
        .unwrap_or_else(|| ts.to_string())
}

// ── HumanReadable impls ──────────────────────────────────────────────

impl HumanReadable for TrainStartResult {
    fn print_human(&self) {
        println!(
            "{} Training job {} started (status: {})",
            "OK".green().bold(),
            self.job_id.bold(),
            colorize_status(&self.status),
        );
        if !self.message.is_empty() {
            println!("  {}", self.message.dimmed());
        }
    }
}

impl HumanReadable for TrainStopResult {
    fn print_human(&self) {
        if self.success {
            println!("{} {}", "OK".green().bold(), self.message);
        } else {
            println!("{} {}", "FAILED".red().bold(), self.message);
        }
    }
}

impl HumanReadable for TrainStatusResult {
    fn print_human(&self) {
        println!("{} {}", "Job:".bold(), self.job_id);
        println!("  Model:       {}", self.model_type);
        println!("  Status:      {}", colorize_status(&self.status));
        if !self.description.is_empty() {
            println!("  Description: {}", self.description);
        }
        println!("  Created:     {}", format_timestamp(self.created_at));
        if self.started_at > 0 {
            println!("  Started:     {}", format_timestamp(self.started_at));
        }
        if self.completed_at > 0 {
            println!("  Completed:   {}", format_timestamp(self.completed_at));
        }
        if !self.error.is_empty() {
            println!("  Error:       {}", self.error.red());
        }
        if !self.artifact_path.is_empty() {
            println!("  Artifact:    {}", self.artifact_path.dimmed());
        }
        if let Some(fm) = &self.financial_metrics {
            println!("\n  {}:", "Financial Metrics".bold());
            println!("    Sharpe:     {:.3}", fm.sharpe_ratio);
            println!("    MaxDD:      {:.2}%", fm.max_drawdown * 100.0);
            println!("    Hit Rate:   {:.1}%", fm.hit_rate * 100.0);
            println!("    Return:     {:.2}%", fm.simulated_return * 100.0);
        }
    }
}

impl HumanReadable for TrainListResult {
    fn print_human(&self) {
        println!("{} training job(s)\n", self.total);
        println!(
            "  {:<20} {:<12} {:<12} {:>10} {:>10}  {}",
            "JOB ID".bold(),
            "MODEL".bold(),
            "STATUS".bold(),
            "LOSS".bold(),
            "BEST VAL".bold(),
            "CREATED".bold(),
        );
        println!("  {}", "-".repeat(86));

        for j in &self.jobs {
            let loss_str = if j.final_loss > 0.0 {
                format!("{:.4}", j.final_loss)
            } else {
                "-".to_owned()
            };
            let val_str = if j.best_val > 0.0 {
                format!("{:.4}", j.best_val)
            } else {
                "-".to_owned()
            };
            println!(
                "  {:<20} {:<12} {:<12} {:>10} {:>10}  {}",
                truncate_id(&j.job_id, 18),
                j.model_type,
                colorize_status(&j.status),
                loss_str,
                val_str,
                format_timestamp(j.created_at).dimmed(),
            );
        }
    }
}

fn truncate_id(id: &str, max: usize) -> String {
    if id.len() <= max {
        id.to_owned()
    } else {
        // Safe: only truncate at ASCII boundaries (UUIDs and job IDs are ASCII)
        let suffix_len = max.saturating_sub(2);
        let chars: Vec<char> = id.chars().collect();
        let start = chars.len().saturating_sub(suffix_len);
        let tail: String = chars.get(start..).unwrap_or_default().iter().collect();
        format!("..{tail}")
    }
}

// ── Execute ──────────────────────────────────────────────────────────

impl TrainCommand {
    pub async fn execute(&self, client: &FoxhuntClient, format: OutputFormat) -> Result<()> {
        match &self.action {
            TrainAction::Start { model, config, gpu } => {
                let data_source = config.as_ref().map(|path| ml_training::DataSource {
                    source: Some(ml_training::data_source::Source::FilePath(path.clone())),
                    start_time: 0,
                    end_time: 0,
                });

                let resp = client
                    .ml_training()
                    .start_training(ml_training::StartTrainingRequest {
                        model_type: model.to_uppercase(),
                        data_source,
                        hyperparameters: None,
                        use_gpu: *gpu,
                        description: String::new(),
                        tags: Default::default(),
                        mode: ml_training::TrainingMode::Full as i32,
                        resume_checkpoint_path: String::new(),
                        max_epochs: 0,
                    })
                    .await?
                    .into_inner();

                let result = TrainStartResult {
                    job_id: resp.job_id,
                    status: training_status_str(resp.status).to_owned(),
                    message: resp.message,
                };
                output::render(&result, format)?;
            }
            TrainAction::Stop { job_id } => {
                let resp = client
                    .ml_training()
                    .stop_training(ml_training::StopTrainingRequest {
                        job_id: job_id.clone(),
                        reason: String::new(),
                    })
                    .await?
                    .into_inner();

                let result = TrainStopResult {
                    success: resp.success,
                    message: resp.message,
                };
                output::render(&result, format)?;
            }
            TrainAction::Status { job_id } => {
                let resp = client
                    .ml_training()
                    .get_training_job_details(
                        ml_training::GetTrainingJobDetailsRequest {
                            job_id: job_id.clone(),
                        },
                    )
                    .await?
                    .into_inner();

                let details = resp.job_details.unwrap_or_default();
                let fm = details.final_financial_metrics.map(|m| {
                    FinancialMetricsSummary {
                        sharpe_ratio: m.sharpe_ratio,
                        max_drawdown: m.max_drawdown,
                        hit_rate: m.hit_rate,
                        simulated_return: m.simulated_return,
                    }
                });

                let result = TrainStatusResult {
                    job_id: details.job_id,
                    model_type: details.model_type,
                    status: training_status_str(details.status).to_owned(),
                    description: details.description,
                    created_at: details.created_at,
                    started_at: details.started_at,
                    completed_at: details.completed_at,
                    error: details.error_message,
                    artifact_path: details.model_artifact_path,
                    financial_metrics: fm,
                };
                output::render(&result, format)?;
            }
            TrainAction::List { status, model } => {
                let status_filter = status
                    .as_deref()
                    .map(status_filter_to_i32)
                    .unwrap_or(0);

                let resp = client
                    .ml_training()
                    .list_training_jobs(ml_training::ListTrainingJobsRequest {
                        page: 0,
                        page_size: 50,
                        status_filter,
                        model_type_filter: model.clone().unwrap_or_default(),
                        start_time: 0,
                        end_time: 0,
                    })
                    .await?
                    .into_inner();

                let jobs: Vec<TrainJobRow> = resp
                    .jobs
                    .iter()
                    .map(|j| TrainJobRow {
                        job_id: j.job_id.clone(),
                        model_type: j.model_type.clone(),
                        status: training_status_str(j.status).to_owned(),
                        description: j.description.clone(),
                        final_loss: j.final_loss,
                        best_val: j.best_validation_score,
                        created_at: j.created_at,
                    })
                    .collect();

                let result = TrainListResult {
                    total: resp.total_count,
                    jobs,
                };
                output::render(&result, format)?;
            }
            TrainAction::Monitor {
                once,
                model,
                interval,
            } => {
                let model_filter = model.clone().unwrap_or_default();

                if *once {
                    let resp = client
                        .monitoring()
                        .get_live_training_metrics(
                            monitoring::GetLiveTrainingMetricsRequest {
                                model_filter,
                            },
                        )
                        .await?
                        .into_inner();
                    render_monitor_snapshot(&resp);
                } else {
                    println!(
                        "{} Streaming training metrics (Ctrl+C to stop)\n",
                        "LIVE".cyan().bold(),
                    );

                    let mut stream = client
                        .monitoring()
                        .stream_training_metrics(
                            monitoring::StreamTrainingMetricsRequest {
                                model_filter,
                                interval_seconds: *interval,
                            },
                        )
                        .await?
                        .into_inner();

                    let mut first = true;
                    while let Some(msg) = stream.next().await {
                        let snapshot = msg?;
                        if !first {
                            // Move cursor up to overwrite previous output.
                            let lines_to_clear =
                                3 + snapshot.sessions.len() + 4 + 1; // header+sessions+gpu+blank
                            print!("\x1b[{}A\x1b[J", lines_to_clear);
                        }
                        first = false;
                        render_monitor_snapshot(&snapshot);
                        io::stdout().flush().ok();
                    }
                }
            }
            TrainAction::Logs { job_id, follow } => {
                if *follow {
                    // Streaming mode: subscribe to live training status updates
                    let mut stream = client
                        .ml_training()
                        .subscribe_to_training_status(
                            ml_training::SubscribeToTrainingStatusRequest {
                                job_id: job_id.clone(),
                            },
                        )
                        .await?
                        .into_inner();

                    println!(
                        "{} Streaming training updates for {} (Ctrl+C to stop)\n",
                        "LIVE".cyan().bold(),
                        job_id.bold(),
                    );

                    while let Some(msg) = stream.next().await {
                        let upd = msg?;
                        let status = colorize_status(training_status_str(upd.status));
                        let ts = format_timestamp(upd.timestamp);

                        println!(
                            "[{}] {} epoch {}/{} progress {:.1}% | {}",
                            ts.dimmed(),
                            status,
                            upd.current_epoch,
                            upd.total_epochs,
                            upd.progress_percentage,
                            upd.message,
                        );

                        if !upd.metrics.is_empty() {
                            let parts: Vec<String> = upd
                                .metrics
                                .iter()
                                .map(|(k, v)| format!("{k}={v:.4}"))
                                .collect();
                            println!("  metrics: {}", parts.join(", ").dimmed());
                        }
                    }
                } else {
                    // Non-follow mode: show job details with status history
                    let resp = client
                        .ml_training()
                        .get_training_job_details(
                            ml_training::GetTrainingJobDetailsRequest {
                                job_id: job_id.clone(),
                            },
                        )
                        .await?
                        .into_inner();

                    let details = resp.job_details.unwrap_or_default();

                    println!("{} {} ({})\n", "Job:".bold(), details.job_id, details.model_type);

                    if details.status_history.is_empty() {
                        println!("  No status updates recorded.");
                    } else {
                        for upd in &details.status_history {
                            let status = colorize_status(training_status_str(upd.status));
                            let ts = format_timestamp(upd.timestamp);
                            println!(
                                "  [{}] {} epoch {}/{} {:.1}% - {}",
                                ts.dimmed(),
                                status,
                                upd.current_epoch,
                                upd.total_epochs,
                                upd.progress_percentage,
                                upd.message,
                            );
                        }
                    }

                    println!(
                        "\n  Tip: use {} to stream live updates.",
                        "--follow".bold()
                    );
                }
            }
        }
        Ok(())
    }
}

// ── Monitor rendering ────────────────────────────────────────────────

fn render_monitor_snapshot(resp: &monitoring::GetLiveTrainingMetricsResponse) {
    let ts = format_timestamp(resp.timestamp);

    // Header
    println!(
        "{} {} | {} K8s job(s) | CPU {:.0}% | RAM {:.0}/{:.0} GB",
        "TRAINING".cyan().bold(),
        ts.dimmed(),
        resp.active_k8s_jobs,
        resp.cpu_percent,
        resp.memory_used_mb / 1024.0,
        resp.memory_total_mb / 1024.0,
    );

    if resp.sessions.is_empty() {
        println!("  {}", "No active training sessions.".dimmed());
        println!();
        return;
    }

    // Sessions table header
    println!(
        "  {:<8} {:<6} {:>5} {:>9} {:>9} {:>7} {:>7} {:>7} {:>7} {:>6}",
        "MODEL".bold(),
        "FOLD".bold(),
        "EPOCH".bold(),
        "LOSS".bold(),
        "VAL_LOSS".bold(),
        "SHARPE".bold(),
        "SORT.".bold(),
        "WIN%".bold(),
        "MAXDD".bold(),
        "BPS".bold(),
    );

    for s in &resp.sessions {
        let model_label = if s.is_hyperopt {
            format!("{}*", s.model)
        } else {
            s.model.clone()
        };

        let loss_str = if s.epoch_loss > 0.0 {
            format!("{:.5}", s.epoch_loss)
        } else {
            "-".to_owned()
        };
        let val_str = if s.validation_loss > 0.0 {
            format!("{:.5}", s.validation_loss)
        } else {
            "-".to_owned()
        };
        let sharpe_str = if s.epoch_sharpe != 0.0 {
            format!("{:.2}", s.epoch_sharpe)
        } else {
            "-".to_owned()
        };
        let sortino_str = if s.epoch_sortino != 0.0 {
            format!("{:.2}", s.epoch_sortino)
        } else {
            "-".to_owned()
        };
        let win_str = if s.epoch_win_rate > 0.0 {
            format!("{:.1}%", s.epoch_win_rate * 100.0)
        } else {
            "-".to_owned()
        };
        let dd_str = if s.epoch_max_drawdown > 0.0 {
            format!("{:.1}%", s.epoch_max_drawdown * 100.0)
        } else {
            "-".to_owned()
        };

        println!(
            "  {:<8} {:<6} {:>5.0} {:>9} {:>9} {:>7} {:>7} {:>7} {:>7} {:>6.1}",
            model_label,
            s.fold,
            s.current_epoch,
            loss_str,
            val_str,
            sharpe_str,
            sortino_str,
            win_str,
            dd_str,
            s.batches_per_second,
        );

        // Show health warnings inline
        if s.nan_detected > 0 || s.gradient_explosions > 0 || s.checkpoint_failures > 0 {
            let mut warnings = Vec::new();
            if s.nan_detected > 0 {
                warnings.push(format!("NaN:{}", s.nan_detected).red().to_string());
            }
            if s.gradient_explosions > 0 {
                warnings.push(format!("grad_exp:{}", s.gradient_explosions).red().to_string());
            }
            if s.checkpoint_failures > 0 {
                warnings.push(format!("ckpt_fail:{}", s.checkpoint_failures).yellow().to_string());
            }
            println!("    {}", warnings.join(" | "));
        }

        // Hyperopt progress
        if s.is_hyperopt && s.hyperopt_trial_total > 0 {
            println!(
                "    {} trial {}/{} best={:.4}",
                "HYPEROPT".magenta(),
                s.hyperopt_trial_current,
                s.hyperopt_trial_total,
                s.hyperopt_best_objective,
            );
        }
    }

    // GPU panel
    if let Some(gpu) = &resp.gpu {
        let temp_color = if gpu.temperature_celsius > 85.0 {
            "red"
        } else if gpu.temperature_celsius > 75.0 {
            "yellow"
        } else {
            "green"
        };
        let temp_str = match temp_color {
            "red" => format!("{:.0}C", gpu.temperature_celsius).red().to_string(),
            "yellow" => format!("{:.0}C", gpu.temperature_celsius).yellow().to_string(),
            _ => format!("{:.0}C", gpu.temperature_celsius).green().to_string(),
        };

        println!(
            "\n  {} util {:.0}% | VRAM {:.1}/{:.0} GB | {} | {:.0}W",
            "GPU".bold(),
            gpu.utilization_percent,
            gpu.memory_used_mb / 1024.0,
            gpu.memory_total_mb / 1024.0,
            temp_str,
            gpu.power_watts,
        );
    }

    println!();
}
