zilliz 1.4.1

use anyhow::{bail, Context, Result};

use crate::api::client::ApiClient;
use crate::api::error::ApiError;
use crate::cli::args::ContextCommands;
use crate::cli::formatter;
use crate::config::context::ClusterContext;
use crate::config::credentials::resolve_api_key;
use crate::config::manager::ConfigManager;
use crate::model::loader::Models;

pub async fn run(
    models: &Models,
    config_mgr: &ConfigManager,
    cmd: ContextCommands,
    output_format: &str,
    api_key_override: Option<&str>,
) -> Result<()> {
    match cmd {
        ContextCommands::Set {
            cluster_id,
            endpoint,
            database,
            on_demand,
        } => {
            set_context(
                models,
                config_mgr,
                cluster_id,
                endpoint,
                database,
                on_demand,
                api_key_override,
            )
            .await
        }
        ContextCommands::Current { output } => {
            let fmt = output.as_deref().unwrap_or(output_format);
            show_current(config_mgr, fmt)
        }
        ContextCommands::Clear => clear_context(config_mgr),
    }
}

async fn set_context(
    models: &Models,
    config_mgr: &ConfigManager,
    cluster_id: Option<String>,
    endpoint: Option<String>,
    database: Option<String>,
    on_demand: bool,
    api_key_override: Option<&str>,
) -> Result<()> {
    // If only database is provided, update just that field
    if let (Some(db), None) = (&database, &cluster_id) {
        let mut ctx = config_mgr.get_context();
        ctx.database = db.clone();
        config_mgr.set_context(&ctx)?;
        println!("Database set to: {}", ctx.database);
        return Ok(());
    }

    let cluster_id = cluster_id.context("--cluster-id is required (or provide only --database)")?;

    // Resolve endpoint from cluster API if not provided
    let (resolved_endpoint, plan) = if let Some(ep) = endpoint {
        (ep, None)
    } else {
        let api_key = resolve_api_key(api_key_override, config_mgr).ok_or(ApiError::NoApiKey)?;
        let base_url =
            super::endpoint::resolve_control_plane_url(config_mgr, &models.control_plane, None);
        let client = ApiClient::new(api_key, base_url);

        let path = if on_demand {
            format!("/v2/clusters/onDemandClusters/{}", cluster_id)
        } else {
            format!("/v2/clusters/{}", cluster_id)
        };
        let result = client.call("GET", &path, None, None).await?;

        let endpoint_field = if on_demand {
            "endpoint"
        } else {
            "connectAddress"
        };
        let ep = result
            .get(endpoint_field)
            .and_then(|v| v.as_str())
            .map(|s| s.to_string())
            .with_context(|| {
                format!(
                    "Cluster has no {} field. Use --endpoint to set manually.",
                    endpoint_field
                )
            })?;

        let plan = result
            .get("plan")
            .or_else(|| result.get("deploymentOption"))
            .and_then(|v| v.as_str())
            .map(|s| s.to_string());

        (ep, plan)
    };

    let db = database.unwrap_or_else(|| "default".to_string());

    config_mgr.clear_context()?;
    config_mgr.set_context(&ClusterContext {
        cluster_id: Some(cluster_id.clone()),
        endpoint: Some(resolved_endpoint.clone()),
        database: db.clone(),
        plan,
    })?;

    println!(
        "Context set: cluster={}, endpoint={}, database={}",
        cluster_id, resolved_endpoint, db
    );
    Ok(())
}

fn show_current(config_mgr: &ConfigManager, output_format: &str) -> Result<()> {
    let ctx = config_mgr.get_context();
    if !ctx.is_set() {
        bail!("No context set. Run `zilliz context set --cluster-id <id>` first.");
    }

    match output_format {
        "json" => {
            let data = serde_json::json!({
                "cluster_id": ctx.cluster_id,
                "endpoint": ctx.endpoint,
                "database": ctx.database,
                "plan": ctx.plan,
            });
            println!("{}", formatter::format_json(&data));
        }
        _ => {
            println!("Cluster:  {}", ctx.cluster_id.unwrap_or_default());
            println!("Endpoint: {}", ctx.endpoint.unwrap_or_default());
            println!("Database: {}", ctx.database);
            if let Some(plan) = &ctx.plan {
                println!("Plan:     {}", plan);
            }
        }
    }
    Ok(())
}

fn clear_context(config_mgr: &ConfigManager) -> Result<()> {
    config_mgr.clear_context()?;
    println!("Context cleared.");
    Ok(())
}