//! `fxt tune` -- hyperparameter optimization.

use anyhow::Result;
use clap::{Parser, Subcommand};
use colored::Colorize;
use serde::Serialize;

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

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

#[derive(Subcommand, Debug)]
enum TuneAction {
    /// Start a hyperparameter optimization run
    Start {
        /// Model type
        #[arg(long)]
        model: String,
        /// Number of trials
        #[arg(long, default_value = "50")]
        trials: u32,
        /// Path to search-space config
        #[arg(long)]
        config: Option<String>,
        /// Enable GPU
        #[arg(long)]
        gpu: bool,
        /// Full 3-phase pipeline: BC → Online RL → Refinement (100 epochs/trial)
        #[arg(long)]
        full: bool,
    },
    /// Stop a running optimization
    Stop {
        /// Optimization job ID
        job_id: String,
    },
    /// Show optimization status and best trial so far
    Status {
        /// Optimization job ID
        job_id: String,
    },
    /// Approve best hyperparameters for training
    Approve {
        /// Optimization job ID
        job_id: String,
    },
}

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

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

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

#[derive(Serialize)]
struct TuneStatusResult {
    job_id: String,
    status: String,
    current_trial: u32,
    total_trials: u32,
    best_params: Vec<(String, f32)>,
    best_metrics: Vec<(String, f32)>,
    message: String,
    started_at: i64,
    updated_at: i64,
    trials: Vec<TrialRow>,
}

#[derive(Serialize)]
struct TrialRow {
    trial: u32,
    objective: f32,
    state: String,
    started_at: i64,
    completed_at: i64,
}

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

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

fn tuning_status_str(s: i32) -> &'static str {
    use crate::proto::ml_training::TuningJobStatus;
    match TuningJobStatus::try_from(s) {
        Ok(TuningJobStatus::TuningPending) => "pending",
        Ok(TuningJobStatus::TuningRunning) => "running",
        Ok(TuningJobStatus::TuningCompleted) => "completed",
        Ok(TuningJobStatus::TuningFailed) => "failed",
        Ok(TuningJobStatus::TuningStopped) => "stopped",
        _ => "unknown",
    }
}

fn trial_state_str(s: i32) -> &'static str {
    use crate::proto::ml_training::TrialState;
    match TrialState::try_from(s) {
        Ok(TrialState::TrialRunning) => "running",
        Ok(TrialState::TrialComplete) => "complete",
        Ok(TrialState::TrialPruned) => "pruned",
        Ok(TrialState::TrialFailed) => "failed",
        _ => "unknown",
    }
}

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

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 TuneStartResult {
    fn print_human(&self) {
        println!(
            "{} Tuning job {} started (status: {})",
            "OK".green().bold(),
            self.job_id.bold(),
            colorize_tuning(&self.status),
        );
        if !self.message.is_empty() {
            println!("  {}", self.message.dimmed());
        }
    }
}

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

impl HumanReadable for TuneStatusResult {
    fn print_human(&self) {
        println!("{} {}", "Job:".bold(), self.job_id);
        println!(
            "  Status:   {} (trial {}/{})",
            colorize_tuning(&self.status),
            self.current_trial,
            self.total_trials,
        );
        println!("  Started:  {}", format_timestamp(self.started_at));
        println!("  Updated:  {}", format_timestamp(self.updated_at));
        if !self.message.is_empty() {
            println!("  Message:  {}", self.message.dimmed());
        }

        if !self.best_params.is_empty() {
            println!("\n  {}:", "Best Parameters".bold());
            for (k, v) in &self.best_params {
                println!("    {k:<30} {v:.6}");
            }
        }

        if !self.best_metrics.is_empty() {
            println!("\n  {}:", "Best Metrics".bold());
            for (k, v) in &self.best_metrics {
                println!("    {k:<30} {v:.6}");
            }
        }

        if !self.trials.is_empty() {
            println!(
                "\n  {}:\n  {:<8} {:>12} {:<12} {}",
                "Trial History".bold(),
                "TRIAL".bold(),
                "OBJECTIVE".bold(),
                "STATE".bold(),
                "COMPLETED".bold(),
            );
            println!("  {}", "-".repeat(56));

            for t in &self.trials {
                let obj_str = if t.objective != 0.0 {
                    format!("{:.6}", t.objective)
                } else {
                    "-".to_owned()
                };
                println!(
                    "  {:<8} {:>12} {:<12} {}",
                    t.trial,
                    obj_str,
                    colorize_tuning(&t.state),
                    format_timestamp(t.completed_at).dimmed(),
                );
            }
        }
    }
}

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

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

impl TuneCommand {
    pub async fn execute(&self, client: &FoxhuntClient, format: OutputFormat) -> Result<()> {
        match &self.action {
            TuneAction::Start {
                model,
                trials,
                config,
                gpu,
                full,
            } => {
                let description = if *full {
                    "Full 3-phase pipeline: BC → Online RL → Refinement".to_owned()
                } else {
                    String::new()
                };
                let num_trials = if *full && *trials == 50 { 20 } else { *trials };
                let resp = client
                    .ml_training()
                    .start_tuning_job(ml_training::StartTuningJobRequest {
                        model_type: model.to_uppercase(),
                        num_trials,
                        config_path: config.clone().unwrap_or_default(),
                        data_source: None,
                        use_gpu: *gpu,
                        description,
                        tags: Default::default(),
                    })
                    .await?
                    .into_inner();

                let result = TuneStartResult {
                    job_id: resp.job_id,
                    status: tuning_status_str(resp.status).to_owned(),
                    message: resp.message,
                };
                output::render(&result, format)?;
            }
            TuneAction::Stop { job_id } => {
                let resp = client
                    .ml_training()
                    .stop_tuning_job(ml_training::StopTuningJobRequest {
                        job_id: job_id.clone(),
                        reason: String::new(),
                    })
                    .await?
                    .into_inner();

                let result = TuneStopResult {
                    success: resp.success,
                    message: resp.message,
                    final_status: tuning_status_str(resp.final_status).to_owned(),
                };
                output::render(&result, format)?;
            }
            TuneAction::Status { job_id } => {
                let resp = client
                    .ml_training()
                    .get_tuning_job_status(ml_training::GetTuningJobStatusRequest {
                        job_id: job_id.clone(),
                    })
                    .await?
                    .into_inner();

                let best_params: Vec<(String, f32)> =
                    resp.best_params.into_iter().collect();
                let best_metrics: Vec<(String, f32)> =
                    resp.best_metrics.into_iter().collect();

                let trials: Vec<TrialRow> = resp
                    .trial_history
                    .iter()
                    .map(|t| TrialRow {
                        trial: t.trial_number,
                        objective: t.objective_value,
                        state: trial_state_str(t.state).to_owned(),
                        started_at: t.started_at,
                        completed_at: t.completed_at,
                    })
                    .collect();

                let result = TuneStatusResult {
                    job_id: resp.job_id,
                    status: tuning_status_str(resp.status).to_owned(),
                    current_trial: resp.current_trial,
                    total_trials: resp.total_trials,
                    best_params,
                    best_metrics,
                    message: resp.message,
                    started_at: resp.started_at,
                    updated_at: resp.updated_at,
                    trials,
                };
                output::render(&result, format)?;
            }
            TuneAction::Approve { job_id } => {
                let resp = client
                    .ml_training()
                    .approve_promotion(ml_training::ApprovePromotionRequest {
                        model_id: job_id.clone(),
                    })
                    .await?
                    .into_inner();

                let result = TuneApproveResult {
                    success: resp.success,
                    message: resp.message,
                };
                output::render(&result, format)?;
            }
        }
        Ok(())
    }
}
