zilliz 1.4.2

use std::collections::HashMap;
use std::io::IsTerminal;

use comfy_table::presets::UTF8_FULL_CONDENSED;
use comfy_table::{Cell, Color, ContentArrangement, Table};
use serde_json::Value;

/// Return terminal width if stdout is a terminal, `None` if piped.
pub fn terminal_width() -> Option<u16> {
    if std::io::stdout().is_terminal() {
        crossterm::terminal::size().ok().map(|(w, _)| w)
    } else {
        None
    }
}

/// Create a `comfy_table::Table` with standard preset and dynamic content
/// arrangement. When stdout is a TTY the table width is capped at the
/// terminal width; piped output is unbounded.
pub fn create_table(columns: &[impl AsRef<str>], no_header: bool) -> Table {
    let mut table = Table::new();
    table.load_preset(UTF8_FULL_CONDENSED);
    table.set_content_arrangement(ContentArrangement::Dynamic);
    if let Some(w) = terminal_width() {
        table.set_width(w);
    }
    if !no_header {
        let headers: Vec<&str> = columns.iter().map(|c| c.as_ref()).collect();
        table.set_header(headers);
    }
    table
}

/// Format a JSON value as pretty-printed JSON string.
pub fn format_json(value: &Value) -> String {
    serde_json::to_string_pretty(value).unwrap_or_else(|_| value.to_string())
}

/// Format a JSON value as YAML.
pub fn format_yaml(value: &Value) -> String {
    serde_yaml::to_string(value)
        .unwrap_or_else(|_| format_json(value))
        .trim_end()
        .to_string()
}

/// Format a JSON value as CSV.
pub fn format_csv(value: &Value, no_header: bool) -> String {
    let items = match value {
        Value::Array(arr) => arr.iter().collect::<Vec<_>>(),
        Value::Object(_) => vec![value],
        other => return other.to_string(),
    };

    if items.is_empty() {
        return String::new();
    }

    let mut wtr = csv::WriterBuilder::new().from_writer(vec![]);

    if let Some(first) = items.first().and_then(|v| v.as_object()) {
        let keys: Vec<&String> = first.keys().collect();
        if !no_header {
            let headers: Vec<&str> = keys.iter().map(|k| k.as_str()).collect();
            let _ = wtr.write_record(&headers);
        }
        for item in &items {
            let row: Vec<String> = keys
                .iter()
                .map(|k| {
                    item.get(k.as_str())
                        .map(|v| match v {
                            Value::String(s) => s.clone(),
                            Value::Null => String::new(),
                            other => other.to_string(),
                        })
                        .unwrap_or_default()
                })
                .collect();
            let _ = wtr.write_record(&row);
        }
    } else {
        if !no_header {
            let _ = wtr.write_record(["value"]);
        }
        for item in &items {
            let val = match item {
                Value::String(s) => s.clone(),
                other => other.to_string(),
            };
            let _ = wtr.write_record([&val]);
        }
    }

    let _ = wtr.flush();
    String::from_utf8(wtr.into_inner().unwrap_or_default())
        .unwrap_or_default()
        .trim_end_matches('\n')
        .to_string()
}

/// Apply a JMESPath query filter to a JSON value.
pub fn apply_query(value: &Value, expression: &str) -> anyhow::Result<Value> {
    let expr = jmespath::compile(expression)
        .map_err(|e| anyhow::anyhow!("Invalid JMESPath expression: {}", e))?;
    let jmes_data = jmespath::Variable::from_json(&value.to_string())
        .map_err(|e| anyhow::anyhow!("Failed to parse data for JMESPath: {}", e))?;
    let result = expr
        .search(jmes_data)
        .map_err(|e| anyhow::anyhow!("JMESPath search failed: {}", e))?;
    let json_str = serde_json::to_string(&*result)
        .map_err(|e| anyhow::anyhow!("Failed to serialize JMESPath result: {}", e))?;
    serde_json::from_str(&json_str)
        .map_err(|e| anyhow::anyhow!("Failed to parse JMESPath result: {}", e))
}

/// Format a JSON value as plain text (key: value pairs for objects, one line per item for arrays).
pub fn format_text(value: &Value) -> String {
    match value {
        Value::Object(map) => map
            .iter()
            .map(|(k, v)| {
                let val = match v {
                    Value::String(s) => s.clone(),
                    Value::Null => "".to_string(),
                    other => other.to_string(),
                };
                format!("{}: {}", k, val)
            })
            .collect::<Vec<_>>()
            .join("\n"),
        Value::Array(arr) => arr
            .iter()
            .map(format_text)
            .collect::<Vec<_>>()
            .join("\n---\n"),
        Value::String(s) => s.clone(),
        Value::Null => "".to_string(),
        other => other.to_string(),
    }
}

/// Parse a color name string to a `comfy_table::Color`.
fn parse_color(name: &str) -> Option<Color> {
    match name.to_lowercase().as_str() {
        "red" => Some(Color::Red),
        "green" => Some(Color::Green),
        "yellow" => Some(Color::Yellow),
        "blue" => Some(Color::Blue),
        "cyan" => Some(Color::Cyan),
        "magenta" => Some(Color::Magenta),
        "white" => Some(Color::White),
        "darkred" => Some(Color::DarkRed),
        "darkgreen" => Some(Color::DarkGreen),
        "darkyellow" => Some(Color::DarkYellow),
        "darkblue" => Some(Color::DarkBlue),
        "darkcyan" => Some(Color::DarkCyan),
        "darkmagenta" => Some(Color::DarkMagenta),
        "grey" | "gray" => Some(Color::Grey),
        _ => None,
    }
}

/// Resolve the color for a cell value given a column's color mapping.
/// Tries the exact value first, then falls back to `"*"`.
fn resolve_cell_color(value: &str, col_colors: &HashMap<String, String>) -> Option<Color> {
    col_colors
        .get(value)
        .or_else(|| col_colors.get("*"))
        .and_then(|name| parse_color(name))
}

/// Format a JSON array of objects as a table.
pub fn format_table(values: &[&Value], columns: &[&str]) -> String {
    format_table_with_opts(values, columns, false, None)
}

/// Format a JSON array of objects as a table, with optional header suppression and color map.
pub fn format_table_with_opts(
    values: &[&Value],
    columns: &[&str],
    no_header: bool,
    color_map: Option<&HashMap<String, HashMap<String, String>>>,
) -> String {
    let is_tty = terminal_width().is_some();
    let mut table = create_table(columns, no_header);

    for item in values {
        let cells: Vec<Cell> = columns
            .iter()
            .map(|col| {
                let val = item
                    .get(col)
                    .map(|v| match v {
                        Value::String(s) => s.clone(),
                        Value::Null => "".to_string(),
                        Value::Array(arr) => format_scalar_array(arr),
                        other => other.to_string(),
                    })
                    .unwrap_or_default();
                let mut cell = Cell::new(&val);
                if is_tty {
                    if let Some(cm) = color_map {
                        if let Some(col_colors) = cm.get(*col) {
                            if let Some(color) = resolve_cell_color(&val, col_colors) {
                                cell = cell.fg(color);
                            }
                        }
                    }
                }
                cell
            })
            .collect();
        table.add_row(cells);
    }

    table.to_string()
}

/// Unwrap a single protobuf-style array value.
///
/// The Milvus REST API returns Array fields in protobuf wrapping:
///   `{"Data": {"StringData": {"data": ["a", "b"]}}}`
/// This function unwraps it to the plain array: `["a", "b"]`.
fn unwrap_protobuf_value(value: &Value) -> Option<Value> {
    let outer = value.as_object()?;
    let data_inner = outer.get("Data")?.as_object()?;
    if data_inner.len() != 1 {
        return None;
    }
    let type_wrapper = data_inner.values().next()?.as_object()?;
    type_wrapper.get("data").cloned()
}

/// Normalize protobuf-wrapped Array fields in API response data.
///
/// Walks each item in the result and replaces protobuf-wrapped values with
/// their unwrapped plain arrays, matching the zilliz-cli behavior.
pub fn normalize_array_fields(value: &mut Value) {
    fn normalize_object(map: &mut serde_json::Map<String, Value>) {
        let updates: Vec<(String, Value)> = map
            .iter()
            .filter_map(|(k, v)| unwrap_protobuf_value(v).map(|nv| (k.clone(), nv)))
            .collect();
        for (k, v) in updates {
            map.insert(k, v);
        }
    }

    match value {
        Value::Array(arr) => {
            for item in arr.iter_mut() {
                if let Value::Object(map) = item {
                    normalize_object(map);
                }
            }
        }
        Value::Object(map) => {
            // First normalize any nested arrays (e.g. {"data": [{...}]})
            for v in map.values_mut() {
                if let Value::Array(arr) = v {
                    for item in arr.iter_mut() {
                        if let Value::Object(inner) = item {
                            normalize_object(inner);
                        }
                    }
                }
            }
            // Then normalize the object itself
            normalize_object(map);
        }
        _ => {}
    }
}

/// Check if a JSON array contains only scalar (non-object, non-array) values.
fn is_scalar_array(arr: &[Value]) -> bool {
    arr.iter()
        .all(|v| !matches!(v, Value::Object(_) | Value::Array(_)))
}

/// Format a JSON array of scalars as a comma-separated string.
fn format_scalar_array(arr: &[Value]) -> String {
    arr.iter()
        .map(|v| match v {
            Value::String(s) => s.clone(),
            Value::Null => String::new(),
            other => other.to_string(),
        })
        .collect::<Vec<_>>()
        .join(", ")
}

/// Auto-detect column names from the first item in a list of JSON objects.
/// Filters out overly nested or large fields. Sorts name/id columns to the front.
pub fn auto_columns(items: &[&Value]) -> Vec<String> {
    let first = match items.first() {
        Some(item) => item,
        None => return vec![],
    };

    let obj = match first.as_object() {
        Some(o) => o,
        None => return vec![],
    };

    let mut cols: Vec<String> = obj
        .keys()
        .filter(|k| {
            match first.get(k.as_str()) {
                // Skip nested objects
                Some(Value::Object(_)) => false,
                // Allow scalar arrays (e.g. [1, 2, 3] or ["a", "b"]),
                // skip arrays of objects (nested tables)
                Some(Value::Array(arr)) => is_scalar_array(arr),
                _ => true,
            }
        })
        .cloned()
        .collect();

    // Sort: name-like columns first, then id-like columns, then the rest
    cols.sort_by_key(|k| {
        let lower = k.to_lowercase();
        if lower.ends_with("name") || lower == "name" {
            0
        } else if lower.ends_with("id") || lower == "id" {
            1
        } else {
            2
        }
    });

    cols
}

/// Look up a value by key from a params array of `{"key": "...", "value": "..."}` pairs.
fn lookup_param(params: &[Value], key: &str) -> Option<String> {
    params.iter().find_map(|entry| {
        let obj = entry.as_object()?;
        let k = obj.get("key").and_then(|v| v.as_str())?;
        if k == key {
            Some(
                obj.get("value")
                    .map(|v| match v {
                        Value::String(s) => s.clone(),
                        Value::Null => String::new(),
                        other => other.to_string(),
                    })
                    .unwrap_or_default(),
            )
        } else {
            None
        }
    })
}

/// Columns to display in the fields sub-table of `collection describe`.
const FIELD_DISPLAY_COLUMNS: &[&str] = &[
    "name",
    "type",
    "primaryKey",
    "autoId",
    "nullable",
    "description",
];

/// Enrich field items for display: embed params info into the type column.
///
/// - Vector types: `FloatVector(128)` (dim from params)
/// - VarChar: `VarChar(256)` (max_length from params)
/// - Array: `Array<Int8>[32]` or `Array<VarChar(256)>[32]`
///   (elementType + max_capacity from params, max_length for VarChar elements)
///
/// Params is an array of `{"key": "...", "value": "..."}` pairs.
///
/// Returns enriched copies if items look like Milvus field descriptors
/// (have both "type" and "name" keys), otherwise returns None.
fn enrich_field_items(items: &[&Value]) -> Option<Vec<Value>> {
    // Only apply if items look like field descriptors
    let first = items.first()?.as_object()?;
    if !first.contains_key("type") || !first.contains_key("name") {
        return None;
    }

    let vector_types = [
        "FloatVector",
        "Float16Vector",
        "BFloat16Vector",
        "BinaryVector",
        "SparseFloatVector",
    ];

    let enriched: Vec<Value> = items
        .iter()
        .map(|item| {
            let mut obj = item.as_object().cloned().unwrap_or_default();
            let type_str = obj
                .get("type")
                .and_then(|v| v.as_str())
                .unwrap_or_default()
                .to_string();
            let params: Vec<Value> = obj
                .get("params")
                .and_then(|v| v.as_array())
                .cloned()
                .unwrap_or_default();

            let new_type = if vector_types
                .iter()
                .any(|t| type_str.eq_ignore_ascii_case(t))
            {
                match lookup_param(&params, "dim") {
                    Some(d) => format!("{}({})", type_str, d),
                    None => type_str,
                }
            } else if type_str.eq_ignore_ascii_case("VarChar") {
                match lookup_param(&params, "max_length") {
                    Some(ml) => format!("{}({})", type_str, ml),
                    None => type_str,
                }
            } else if type_str.eq_ignore_ascii_case("Array") {
                let elem_type = obj
                    .get("elementType")
                    .and_then(|v| v.as_str())
                    .unwrap_or("Unknown");
                let max_cap = lookup_param(&params, "max_capacity");
                let elem_display = if elem_type.eq_ignore_ascii_case("VarChar") {
                    match lookup_param(&params, "max_length") {
                        Some(ml) => format!("{}({})", elem_type, ml),
                        None => elem_type.to_string(),
                    }
                } else {
                    elem_type.to_string()
                };
                match max_cap {
                    Some(mc) => format!("Array<{}>[{}]", elem_display, mc),
                    None => format!("Array<{}>", elem_display),
                }
            } else {
                type_str
            };

            obj.insert("type".to_string(), Value::String(new_type));
            // Remove params and elementType since they are now embedded in type
            obj.remove("params");
            obj.remove("elementType");
            Value::Object(obj)
        })
        .collect();

    Some(enriched)
}

/// Format a single JSON object as a vertical key-value table (key on left, value on right).
/// Nested array-of-objects and object fields are rendered as separate labeled sub-tables.
pub fn format_kv_table(value: &Value) -> String {
    let obj = match value.as_object() {
        Some(o) => o,
        None => return format_json(value),
    };

    // Separate scalar fields from nested fields (arrays-of-objects, objects)
    let mut scalar_rows: Vec<(String, String)> = Vec::new();
    // Preserve original key order for nested sub-tables
    let mut nested_sections: Vec<(String, String)> = Vec::new();

    for (k, v) in obj.iter() {
        match v {
            Value::Array(arr) if !arr.is_empty() && arr.iter().all(|item| item.is_object()) => {
                // Array of objects: render as a separate horizontal table
                let items: Vec<&Value> = arr.iter().collect();
                // Try to enrich field items (embeds params into type display)
                let sub_table = if let Some(enriched) = enrich_field_items(&items) {
                    // Filter to curated columns, filling missing keys with empty strings
                    let filtered: Vec<Value> = enriched
                        .iter()
                        .map(|item| {
                            let obj = item.as_object();
                            let mut new_obj = serde_json::Map::new();
                            for &col in FIELD_DISPLAY_COLUMNS {
                                let val = obj
                                    .and_then(|o| o.get(col))
                                    .cloned()
                                    .unwrap_or(Value::String(String::new()));
                                new_obj.insert(col.to_string(), val);
                            }
                            Value::Object(new_obj)
                        })
                        .collect();
                    let refs: Vec<&Value> = filtered.iter().collect();
                    format_table_with_opts(&refs, FIELD_DISPLAY_COLUMNS, false, None)
                } else {
                    let auto_cols = auto_columns(&items);
                    let col_refs: Vec<&str> = auto_cols.iter().map(|s| s.as_str()).collect();
                    format_table_with_opts(&items, &col_refs, false, None)
                };
                nested_sections.push((k.clone(), sub_table));
            }
            Value::Array(arr) if arr.is_empty() => {
                // Empty array: skip entirely
            }
            Value::Array(arr) => {
                // Array of scalars: keep inline
                let items: Vec<String> = arr
                    .iter()
                    .map(|item| match item {
                        Value::String(s) => s.clone(),
                        other => other.to_string(),
                    })
                    .collect();
                scalar_rows.push((k.clone(), items.join(", ")));
            }
            Value::Object(map) if !map.is_empty() => {
                // Non-empty object: render as a separate key-value sub-table
                let mut sub_table = create_table(&["key", "value"], false);
                for (sub_k, sub_v) in map.iter() {
                    let val = match sub_v {
                        Value::String(s) => s.clone(),
                        Value::Null => String::new(),
                        other => other.to_string(),
                    };
                    sub_table.add_row([sub_k.as_str(), val.as_str()]);
                }
                nested_sections.push((k.clone(), sub_table.to_string()));
            }
            Value::Object(_) => {
                // Empty object: skip entirely
            }
            Value::String(s) => {
                scalar_rows.push((k.clone(), s.clone()));
            }
            Value::Null => {
                scalar_rows.push((k.clone(), String::new()));
            }
            other => {
                scalar_rows.push((k.clone(), other.to_string()));
            }
        }
    }

    let mut output = String::new();

    // Main key-value table (scalar fields only)
    if !scalar_rows.is_empty() {
        let mut table = create_table(&["Key", "Value"], false);
        for (k, v) in &scalar_rows {
            table.add_row([k.as_str(), v.as_str()]);
        }
        output.push_str(&table.to_string());
    }

    // Sub-tables for nested fields
    for (label, sub_table) in &nested_sections {
        if !output.is_empty() {
            output.push('\n');
        }
        output.push_str(&format!("\n{}:\n", label));
        output.push_str(sub_table);
    }

    output
}

/// Format an API error (with code) for the given output format.
pub fn format_error(format: &str, code: i64, message: &str) -> String {
    match format {
        "json" => {
            let obj = serde_json::json!({"code": code, "message": message});
            serde_json::to_string_pretty(&obj).unwrap_or_else(|_| obj.to_string())
        }
        _ => format!("Error [{}]: {}", code, message),
    }
}

/// Format a non-API error (no code) for the given output format.
pub fn format_error_simple(format: &str, message: &str) -> String {
    match format {
        "json" => {
            let obj = serde_json::json!({"code": 0, "message": message});
            serde_json::to_string_pretty(&obj).unwrap_or_else(|_| obj.to_string())
        }
        _ => format!("Error: {}", message),
    }
}