foxhunt/bin/fxt/src/commands/model.rs

//! `fxt model` -- model management and promotion.

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;

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

#[derive(Subcommand, Debug)]
enum ModelAction {
    /// List registered models
    List {
        /// Filter by model type
        #[arg(long)]
        model: Option<String>,
        /// Filter by status (staging, production, archived)
        #[arg(long)]
        status: Option<String>,
    },
    /// Show detailed model status
    Status {
        /// Model ID
        model_id: String,
    },
    /// Run inference on a model
    Predict {
        /// Model ID
        model_id: String,
        /// Symbol to predict
        #[arg(long)]
        symbol: String,
    },
    /// Show ensemble composition and weights
    Ensemble {
        /// Symbol for ensemble vote
        #[arg(long, default_value = "ES.FUT")]
        symbol: String,
    },
    /// Promote a model to production
    Promote {
        /// Model ID
        model_id: String,
        /// Reason for promotion
        #[arg(long)]
        reason: Option<String>,
    },
}

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

#[derive(Serialize)]
struct ModelListResult {
    count: usize,
    models: Vec<ModelRow>,
}

#[derive(Serialize)]
struct ModelRow {
    name: String,
    model_type: String,
    description: String,
    symbols: Vec<String>,
}

#[derive(Serialize)]
struct ModelStatusResult {
    name: String,
    state: String,
    health: String,
    error: String,
    last_updated: i64,
    last_prediction: i64,
    metadata: Vec<(String, String)>,
}

#[derive(Serialize)]
struct PredictResult {
    model: String,
    symbol: String,
    prediction_type: String,
    value: f64,
    confidence: f64,
    horizon_minutes: i32,
}

#[derive(Serialize)]
struct EnsembleResult {
    symbol: String,
    consensus: String,
    confidence: f64,
    signal_strength: String,
    votes_buy: i32,
    votes_sell: i32,
    votes_hold: i32,
    total_models: i32,
    individual_votes: Vec<VoteRow>,
}

#[derive(Serialize)]
struct VoteRow {
    model: String,
    prediction: String,
    confidence: f64,
    weight: f64,
}

#[derive(Serialize)]
struct StubMessage {
    message: String,
}

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

fn model_state_str(s: i32) -> &'static str {
    use crate::proto::ml::ModelState;
    match ModelState::try_from(s) {
        Ok(ModelState::Loading) => "loading",
        Ok(ModelState::Ready) => "ready",
        Ok(ModelState::Predicting) => "predicting",
        Ok(ModelState::Training) => "training",
        Ok(ModelState::Error) => "error",
        Ok(ModelState::Offline) => "offline",
        _ => "unknown",
    }
}

fn model_health_str(h: i32) -> &'static str {
    use crate::proto::ml::ModelHealth;
    match ModelHealth::try_from(h) {
        Ok(ModelHealth::Healthy) => "healthy",
        Ok(ModelHealth::Degraded) => "degraded",
        Ok(ModelHealth::Unhealthy) => "unhealthy",
        Ok(ModelHealth::Critical) => "critical",
        _ => "unknown",
    }
}

fn prediction_type_str(p: i32) -> &'static str {
    use crate::proto::ml::PredictionType;
    match PredictionType::try_from(p) {
        Ok(PredictionType::Buy) => "BUY",
        Ok(PredictionType::Sell) => "SELL",
        Ok(PredictionType::Hold) => "HOLD",
        Ok(PredictionType::PriceUp) => "PRICE_UP",
        Ok(PredictionType::PriceDown) => "PRICE_DOWN",
        Ok(PredictionType::VolatilityHigh) => "VOL_HIGH",
        Ok(PredictionType::VolatilityLow) => "VOL_LOW",
        _ => "UNKNOWN",
    }
}

fn signal_strength_str(s: i32) -> &'static str {
    use crate::proto::ml::SignalStrength;
    match SignalStrength::try_from(s) {
        Ok(SignalStrength::VeryWeak) => "very_weak",
        Ok(SignalStrength::Weak) => "weak",
        Ok(SignalStrength::Moderate) => "moderate",
        Ok(SignalStrength::Strong) => "strong",
        Ok(SignalStrength::VeryStrong) => "very_strong",
        _ => "unknown",
    }
}

fn colorize_state(s: &str) -> String {
    match s {
        "ready" | "predicting" => s.green().to_string(),
        "loading" | "training" => s.cyan().to_string(),
        "error" => s.red().bold().to_string(),
        "offline" => s.dimmed().to_string(),
        _ => s.dimmed().to_string(),
    }
}

fn colorize_health(s: &str) -> String {
    match s {
        "healthy" => s.green().bold().to_string(),
        "degraded" => s.yellow().bold().to_string(),
        "unhealthy" => s.red().to_string(),
        "critical" => s.red().bold().to_string(),
        _ => s.dimmed().to_string(),
    }
}

fn colorize_prediction(s: &str) -> String {
    match s {
        "BUY" | "PRICE_UP" => s.green().bold().to_string(),
        "SELL" | "PRICE_DOWN" => s.red().bold().to_string(),
        "HOLD" => s.yellow().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 ModelListResult {
    fn print_human(&self) {
        println!("{} model(s)\n", self.count);
        println!(
            "  {:<24} {:<14} {:<20} {}",
            "NAME".bold(),
            "TYPE".bold(),
            "SYMBOLS".bold(),
            "DESCRIPTION".bold(),
        );
        println!("  {}", "-".repeat(76));

        for m in &self.models {
            let syms = if m.symbols.is_empty() {
                "-".to_owned()
            } else if m.symbols.len() > 3 {
                let first_three = m.symbols.get(..3).map_or_else(
                    || m.symbols.join(", "),
                    |s| s.join(", "),
                );
                format!("{first_three}, +{}", m.symbols.len() - 3)
            } else {
                m.symbols.join(", ")
            };
            println!(
                "  {:<24} {:<14} {:<20} {}",
                m.name,
                m.model_type,
                syms,
                m.description.dimmed(),
            );
        }
    }
}

impl HumanReadable for ModelStatusResult {
    fn print_human(&self) {
        println!("{} {}", "Model:".bold(), self.name);
        println!("  State:          {}", colorize_state(&self.state));
        println!("  Health:         {}", colorize_health(&self.health));
        if !self.error.is_empty() {
            println!("  Error:          {}", self.error.red());
        }
        println!(
            "  Last Updated:   {}",
            format_timestamp(self.last_updated)
        );
        println!(
            "  Last Prediction: {}",
            format_timestamp(self.last_prediction)
        );

        if !self.metadata.is_empty() {
            println!("\n  {}:", "Metadata".bold());
            for (k, v) in &self.metadata {
                println!("    {k}: {v}");
            }
        }
    }
}

impl HumanReadable for PredictResult {
    fn print_human(&self) {
        println!(
            "{} {} on {} -> {} (confidence {:.1}%)",
            "Prediction:".bold(),
            self.model,
            self.symbol,
            colorize_prediction(&self.prediction_type),
            self.confidence * 100.0,
        );
        println!("  Value:    {:.6}", self.value);
        println!("  Horizon:  {} min", self.horizon_minutes);
    }
}

impl HumanReadable for EnsembleResult {
    fn print_human(&self) {
        println!(
            "{} {} -- {} (confidence {:.1}%, signal {})",
            "Ensemble:".bold(),
            self.symbol,
            colorize_prediction(&self.consensus),
            self.confidence * 100.0,
            self.signal_strength,
        );
        println!(
            "  Votes: {} buy / {} sell / {} hold ({} models)",
            self.votes_buy.to_string().green(),
            self.votes_sell.to_string().red(),
            self.votes_hold.to_string().yellow(),
            self.total_models,
        );

        if !self.individual_votes.is_empty() {
            println!(
                "\n  {:<24} {:<12} {:>12} {:>8}",
                "MODEL".bold(),
                "VOTE".bold(),
                "CONFIDENCE".bold(),
                "WEIGHT".bold(),
            );
            println!("  {}", "-".repeat(60));

            for v in &self.individual_votes {
                println!(
                    "  {:<24} {:<12} {:>11.1}% {:>7.3}",
                    v.model,
                    colorize_prediction(&v.prediction),
                    v.confidence * 100.0,
                    v.weight,
                );
            }
        }
    }
}

impl HumanReadable for StubMessage {
    fn print_human(&self) {
        println!("{}", self.message);
    }
}

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

impl ModelCommand {
    pub async fn execute(&self, client: &FoxhuntClient, format: OutputFormat) -> Result<()> {
        match &self.action {
            ModelAction::List { model, status: _ } => {
                let resp = client
                    .ml()
                    .get_available_models(ml::GetAvailableModelsRequest {})
                    .await?
                    .into_inner();

                let mut models: Vec<ModelRow> = resp
                    .available_models
                    .iter()
                    .map(|m| ModelRow {
                        name: m.model_name.clone(),
                        model_type: m.model_type.clone(),
                        description: m.description.clone(),
                        symbols: m.supported_symbols.clone(),
                    })
                    .collect();

                // Apply client-side model type filter if provided.
                if let Some(filter) = model {
                    let filter_upper = filter.to_uppercase();
                    models.retain(|m| m.model_type.to_uppercase().contains(&filter_upper));
                }

                let result = ModelListResult {
                    count: models.len(),
                    models,
                };
                output::render(&result, format)?;
            }
            ModelAction::Status { model_id } => {
                let resp = client
                    .ml()
                    .get_model_status(ml::GetModelStatusRequest {
                        model_name: Some(model_id.clone()),
                    })
                    .await?
                    .into_inner();

                let status = resp
                    .model_statuses
                    .into_iter()
                    .find(|s| s.model_name == *model_id);

                let Some(s) = status else {
                    anyhow::bail!("model '{}' not found", model_id);
                };

                let metadata: Vec<(String, String)> =
                    s.metadata.into_iter().collect();

                let result = ModelStatusResult {
                    name: s.model_name,
                    state: model_state_str(s.state).to_owned(),
                    health: model_health_str(s.health).to_owned(),
                    error: s.error_message.unwrap_or_default(),
                    last_updated: s.last_updated,
                    last_prediction: s.last_prediction,
                    metadata,
                };
                output::render(&result, format)?;
            }
            ModelAction::Predict { model_id, symbol } => {
                let resp = client
                    .ml()
                    .get_prediction(ml::GetPredictionRequest {
                        model_name: model_id.clone(),
                        symbol: symbol.clone(),
                        horizon_minutes: None,
                        features: Default::default(),
                    })
                    .await?
                    .into_inner();

                let pred = resp.prediction.unwrap_or_default();
                let result = PredictResult {
                    model: pred.model_name,
                    symbol: pred.symbol,
                    prediction_type: prediction_type_str(pred.prediction_type).to_owned(),
                    value: pred.value,
                    confidence: resp.confidence,
                    horizon_minutes: pred.horizon_minutes,
                };
                output::render(&result, format)?;
            }
            ModelAction::Ensemble { symbol } => {
                let resp = client
                    .ml()
                    .get_ensemble_vote(ml::GetEnsembleVoteRequest {
                        symbol: symbol.clone(),
                        horizon_minutes: None,
                        model_names: vec![],
                    })
                    .await?
                    .into_inner();

                let vote = resp.ensemble_vote.unwrap_or_default();
                let individual_votes: Vec<VoteRow> = resp
                    .individual_votes
                    .iter()
                    .map(|v| VoteRow {
                        model: v.model_name.clone(),
                        prediction: prediction_type_str(v.prediction).to_owned(),
                        confidence: v.confidence,
                        weight: v.weight,
                    })
                    .collect();

                let result = EnsembleResult {
                    symbol: vote.symbol,
                    consensus: prediction_type_str(vote.consensus_prediction).to_owned(),
                    confidence: resp.overall_confidence,
                    signal_strength: signal_strength_str(vote.signal_strength).to_owned(),
                    votes_buy: vote.votes_buy,
                    votes_sell: vote.votes_sell,
                    votes_hold: vote.votes_hold,
                    total_models: vote.total_models,
                    individual_votes,
                };
                output::render(&result, format)?;
            }
            ModelAction::Promote {
                model_id: _,
                reason: _,
            } => {
                let result = StubMessage {
                    message: "Model promotion is managed through the tuning pipeline.\n\
                              Use `fxt tune approve <job_id>` to promote a tuned model."
                        .to_owned(),
                };
                output::render(&result, format)?;
            }
        }
        Ok(())
    }
}