//! Context window optimization engine.
//!
//! Compresses schema definitions into progressively smaller representations
//! for AI agent context windows, using a five-level progressive disclosure model.

use sha2::{Digest, Sha256};
use std::collections::HashSet;
use std::fmt::Write;

use super::types::*;
use crate::polyglot::extractor::extract_tables;
use crate::schema::types::{SchemaDefinition, TableDef};

/// Approximate token count: ~1 token per 4 characters.
fn estimate_tokens(text: &str) -> usize {
    text.len().div_ceil(4)
}

/// Compute a SHA-256 fingerprint of the schema.
///
/// The fingerprint is a deterministic hash of the schema's serialized JSON
/// representation (with sorted keys). Use this for cache invalidation:
/// if the fingerprint hasn't changed, the schema hasn't changed.
pub fn schema_fingerprint(schema: &SchemaDefinition) -> String {
    // Use JSON with sorted keys for deterministic output
    let json = serde_json::to_string(schema).unwrap_or_default();

    // Sort the JSON keys by serializing to Value first, then back
    let value: serde_json::Value = serde_json::from_str(&json).unwrap_or_default();
    let sorted = serde_json::to_string(&value).unwrap_or_default();

    let mut hasher = Sha256::new();
    hasher.update(sorted.as_bytes());
    let result = hasher.finalize();
    format!("{result:x}")
}

/// Generate a natural language summary of the schema.
///
/// Uses template-based generation (NOT LLM) to produce a human-readable
/// overview of the schema's structure, tables, and relationships.
pub fn schema_summary(schema: &SchemaDefinition) -> SchemaSummary {
    let mut table_summaries = Vec::new();
    let mut table_names: Vec<&String> = schema.tables.keys().collect();
    table_names.sort();

    for name in &table_names {
        let table_def = &schema.tables[*name];
        let description = table_def
            .description
            .clone()
            .unwrap_or_else(|| format!("Table with {} columns", table_def.columns.len()));

        let mut key_columns = Vec::new();
        if let Some(pk) = &table_def.primary_key {
            key_columns.extend(pk.clone());
        }
        // Include FK columns as key columns too
        for fk in &table_def.foreign_keys {
            for col in &fk.columns {
                if !key_columns.contains(col) {
                    key_columns.push(col.clone());
                }
            }
        }

        let relationships: Vec<String> = table_def
            .foreign_keys
            .iter()
            .map(|fk| {
                let from_cols = fk.columns.join(", ");
                let to_cols = fk.references.columns.join(", ");
                format!(
                    "{}.{} -> {}.{}",
                    name, from_cols, fk.references.table, to_cols
                )
            })
            .collect();

        table_summaries.push(TableSummary {
            name: (*name).clone(),
            description,
            column_count: table_def.columns.len(),
            key_columns,
            relationships,
        });
    }

    // Build natural language summary
    let n = table_summaries.len();
    let domain = infer_domain(schema);
    let mut nl = format!(
        "This schema contains {n} table{} ",
        if n == 1 { "" } else { "s" }
    );
    if !domain.is_empty() {
        let _ = write!(nl, "for {domain} ");
    }
    nl.push('(');
    nl.push_str(&schema.dialect.to_string());
    nl.push_str(" dialect):\n");

    for ts in &table_summaries {
        let _ = writeln!(
            nl,
            "- {} ({} columns): {}",
            ts.name, ts.column_count, ts.description
        );
    }

    // Collect all relationships
    let all_rels: Vec<&str> = table_summaries
        .iter()
        .flat_map(|ts| ts.relationships.iter().map(|r| r.as_str()))
        .collect();
    if !all_rels.is_empty() {
        nl.push_str("Key relationships: ");
        nl.push_str(&all_rels.join(", "));
        nl.push('\n');
    }

    let estimated_tokens = estimate_tokens(&nl);

    SchemaSummary {
        natural_language: nl,
        table_summaries,
        estimated_tokens,
    }
}

/// Infer the domain of a schema from table names and descriptions.
fn infer_domain(schema: &SchemaDefinition) -> String {
    let table_names: Vec<String> = schema.tables.keys().map(|n| n.to_lowercase()).collect();

    let all_descriptions: String = schema
        .tables
        .values()
        .filter_map(|t| t.description.as_deref())
        .collect::<Vec<_>>()
        .join(" ")
        .to_lowercase();

    let combined = format!("{} {}", table_names.join(" "), all_descriptions);

    // Simple keyword-based domain inference
    let domains = [
        (
            &["order", "product", "cart", "customer", "payment", "invoice"][..],
            "e-commerce",
        ),
        (
            &["user", "account", "session", "auth", "login", "role"][..],
            "user management",
        ),
        (
            &["campaign", "ad", "impression", "click", "conversion"][..],
            "marketing/advertising",
        ),
        (
            &[
                "patient",
                "doctor",
                "appointment",
                "prescription",
                "diagnosis",
            ][..],
            "healthcare",
        ),
        (
            &["transaction", "balance", "account", "ledger", "transfer"][..],
            "financial",
        ),
        (
            &["event", "metric", "log", "trace", "span"][..],
            "analytics/observability",
        ),
        (
            &["employee", "department", "salary", "payroll"][..],
            "HR/workforce",
        ),
        (
            &["warehouse", "inventory", "shipment", "supplier"][..],
            "supply chain/logistics",
        ),
    ];

    let mut best_domain = "";
    let mut best_score = 0;

    for (keywords, domain) in &domains {
        let score = keywords.iter().filter(|kw| combined.contains(**kw)).count();
        if score > best_score {
            best_score = score;
            best_domain = domain;
        }
    }

    if best_score >= 2 {
        best_domain.to_string()
    } else {
        String::new()
    }
}

/// Optimize a schema for context window usage.
///
/// Compresses the full schema to fit within the configured token budget
/// at the specified disclosure level. If the output exceeds the budget,
/// the engine automatically steps down to a less detailed level.
///
/// # Arguments
/// * `schema` - The full schema definition
/// * `config` - Optimization configuration (token budget, disclosure level, etc.)
///
/// # Returns
/// A `ContextResult` containing the compressed schema and metadata.
pub fn optimize_context(
    schema: &SchemaDefinition,
    config: &ContextConfig,
) -> Result<ContextResult, String> {
    let fingerprint = schema_fingerprint(schema);
    let tables_total = schema.tables.len();

    // Try the requested level, then step down if it exceeds the budget
    let mut level = config.level;
    loop {
        let compressed = render_at_level(schema, level, config, None)?;
        let estimated_tokens = estimate_tokens(&compressed);

        if estimated_tokens <= config.max_tokens || level == DisclosureLevel::Fingerprint {
            let tables_included = count_tables_in_output(&compressed, schema);
            let full_schema_yaml = serde_yaml::to_string(schema).unwrap_or_default();
            let full_tokens = estimate_tokens(&full_schema_yaml);
            let compression_ratio = if full_tokens > 0 {
                estimated_tokens as f64 / full_tokens as f64
            } else {
                1.0
            };

            return Ok(ContextResult {
                compressed_schema: compressed,
                level_used: level,
                estimated_tokens,
                tables_included,
                tables_total,
                compression_ratio,
                fingerprint,
            });
        }

        // Step down to a less detailed level
        match level.prev() {
            Some(prev) => level = prev,
            None => {
                // Should not happen since we check for Fingerprint above
                let compressed =
                    render_at_level(schema, DisclosureLevel::Fingerprint, config, None)?;
                let estimated_tokens = estimate_tokens(&compressed);
                return Ok(ContextResult {
                    compressed_schema: compressed,
                    level_used: DisclosureLevel::Fingerprint,
                    estimated_tokens,
                    tables_included: 0,
                    tables_total,
                    compression_ratio: 0.0,
                    fingerprint,
                });
            }
        }
    }
}

/// Optimize a schema for a specific SQL query's context window.
///
/// Like `optimize_context`, but additionally prunes the schema to only
/// the tables referenced by the query. Relevant tables get full detail
/// while others get reduced detail.
///
/// # Arguments
/// * `sql` - The SQL query to optimize for
/// * `schema` - The full schema definition
/// * `config` - Optimization configuration
///
/// # Returns
/// A `ContextResult` containing the compressed, query-relevant schema.
pub fn optimize_for_query(
    sql: &str,
    schema: &SchemaDefinition,
    config: &ContextConfig,
) -> Result<ContextResult, String> {
    let fingerprint = schema_fingerprint(schema);
    let tables_total = schema.tables.len();

    // Extract tables referenced by the query
    let sql_tables = extract_tables(sql, schema.dialect).unwrap_or_default();

    // Match SQL table names to schema tables (case-insensitive)
    let mut relevant: HashSet<String> = HashSet::new();
    for sql_table in &sql_tables {
        let lower = sql_table.to_lowercase();
        for schema_table in schema.tables.keys() {
            if schema_table.to_lowercase() == lower {
                relevant.insert(schema_table.clone());
            }
        }
    }

    // Also include priority tables
    for pt in &config.priority_tables {
        let lower = pt.to_lowercase();
        for schema_table in schema.tables.keys() {
            if schema_table.to_lowercase() == lower {
                relevant.insert(schema_table.clone());
            }
        }
    }

    // Include FK-referenced tables from relevant tables (one level deep)
    let fk_tables: Vec<String> = relevant
        .iter()
        .filter_map(|t| schema.tables.get(t))
        .flat_map(|td| td.foreign_keys.iter().map(|fk| fk.references.table.clone()))
        .filter(|t| schema.tables.contains_key(t))
        .collect();
    relevant.extend(fk_tables);

    // Try the requested level, then step down if needed
    let mut level = config.level;
    loop {
        let compressed = render_at_level(schema, level, config, Some(&relevant))?;
        let estimated_tokens = estimate_tokens(&compressed);

        if estimated_tokens <= config.max_tokens || level == DisclosureLevel::Fingerprint {
            let tables_included = count_tables_in_output(&compressed, schema);
            let full_schema_yaml = serde_yaml::to_string(schema).unwrap_or_default();
            let full_tokens = estimate_tokens(&full_schema_yaml);
            let compression_ratio = if full_tokens > 0 {
                estimated_tokens as f64 / full_tokens as f64
            } else {
                1.0
            };

            return Ok(ContextResult {
                compressed_schema: compressed,
                level_used: level,
                estimated_tokens,
                tables_included,
                tables_total,
                compression_ratio,
                fingerprint,
            });
        }

        match level.prev() {
            Some(prev) => level = prev,
            None => {
                let compressed =
                    render_at_level(schema, DisclosureLevel::Fingerprint, config, None)?;
                let estimated_tokens = estimate_tokens(&compressed);
                return Ok(ContextResult {
                    compressed_schema: compressed,
                    level_used: DisclosureLevel::Fingerprint,
                    estimated_tokens,
                    tables_included: 0,
                    tables_total,
                    compression_ratio: 0.0,
                    fingerprint,
                });
            }
        }
    }
}

/// Render the schema at a given disclosure level.
fn render_at_level(
    schema: &SchemaDefinition,
    level: DisclosureLevel,
    config: &ContextConfig,
    relevant_tables: Option<&HashSet<String>>,
) -> Result<String, String> {
    let mut table_names: Vec<&String> = schema.tables.keys().collect();
    table_names.sort();

    match level {
        DisclosureLevel::Fingerprint => {
            let fp = schema_fingerprint(schema);
            Ok(format!(
                "schema_fingerprint: {fp}\ntables: {}\n",
                schema.tables.len()
            ))
        }
        DisclosureLevel::TableList => {
            let mut out = String::new();
            let _ = writeln!(out, "# Schema ({} tables)", schema.tables.len());
            for name in &table_names {
                let table_def = &schema.tables[*name];
                let desc = if config.include_descriptions {
                    table_def
                        .description
                        .as_deref()
                        .map(|d| format!(" - {d}"))
                        .unwrap_or_default()
                } else {
                    String::new()
                };
                let _ = writeln!(out, "- {} ({} cols){}", name, table_def.columns.len(), desc);
            }
            Ok(out)
        }
        DisclosureLevel::ColumnNames => {
            let mut out = String::new();
            let _ = writeln!(out, "# Schema ({} tables)", schema.tables.len());
            for name in &table_names {
                let table_def = &schema.tables[*name];
                let desc = if config.include_descriptions {
                    table_def
                        .description
                        .as_deref()
                        .map(|d| format!("  # {d}"))
                        .unwrap_or_default()
                } else {
                    String::new()
                };
                let _ = writeln!(out, "{name}:{desc}");
                let col_names: Vec<&str> =
                    table_def.columns.iter().map(|c| c.name.as_str()).collect();
                let _ = writeln!(out, "  columns: [{}]", col_names.join(", "));

                if config.include_foreign_keys && !table_def.foreign_keys.is_empty() {
                    for fk in &table_def.foreign_keys {
                        let _ = writeln!(
                            out,
                            "  fk: {}.{} -> {}.{}",
                            name,
                            fk.columns.join(", "),
                            fk.references.table,
                            fk.references.columns.join(", ")
                        );
                    }
                }
            }
            Ok(out)
        }
        DisclosureLevel::FullRelevant => {
            let mut out = String::new();
            let _ = writeln!(out, "# Schema ({} tables)", schema.tables.len());

            // Determine which tables are "relevant"
            let priority_lower: HashSet<String> = config
                .priority_tables
                .iter()
                .map(|p| p.to_lowercase())
                .collect();

            for name in &table_names {
                let table_def = &schema.tables[*name];
                let is_relevant = relevant_tables
                    .map(|rt| rt.contains(*name))
                    .unwrap_or(false)
                    || priority_lower.contains(&name.to_lowercase());

                if is_relevant {
                    // Full detail
                    render_table_full(&mut out, name, table_def, config);
                } else {
                    // Compact: just name + column names
                    let col_names: Vec<&str> =
                        table_def.columns.iter().map(|c| c.name.as_str()).collect();
                    let _ = writeln!(out, "{}: [{}]", name, col_names.join(", "));
                }
            }
            Ok(out)
        }
        DisclosureLevel::Full => {
            let mut out = String::new();
            let _ = writeln!(out, "# Schema ({} tables)", schema.tables.len());
            for name in &table_names {
                let table_def = &schema.tables[*name];
                render_table_full(&mut out, name, table_def, config);
            }
            Ok(out)
        }
    }
}

/// Render a single table with full detail.
fn render_table_full(out: &mut String, name: &str, table_def: &TableDef, config: &ContextConfig) {
    let desc = if config.include_descriptions {
        table_def
            .description
            .as_deref()
            .map(|d| format!("  # {d}"))
            .unwrap_or_default()
    } else {
        String::new()
    };
    let _ = writeln!(out, "{name}:{desc}");

    if let Some(pk) = &table_def.primary_key {
        let _ = writeln!(out, "  primary_key: [{}]", pk.join(", "));
    }

    for col in &table_def.columns {
        let nullable = if col.nullable { "?" } else { "" };
        let tags_str = if config.include_tags && !col.tags.is_empty() {
            format!(" [{}]", col.tags.join(", "))
        } else {
            String::new()
        };
        let col_desc = if config.include_descriptions {
            col.description
                .as_deref()
                .map(|d| format!("  # {d}"))
                .unwrap_or_default()
        } else {
            String::new()
        };
        let _ = writeln!(
            out,
            "  - {}: {}{}{}{}",
            col.name, col.data_type, nullable, tags_str, col_desc
        );
    }

    if config.include_foreign_keys && !table_def.foreign_keys.is_empty() {
        let _ = writeln!(out, "  foreign_keys:");
        for fk in &table_def.foreign_keys {
            let _ = writeln!(
                out,
                "    {}.{} -> {}.{}",
                name,
                fk.columns.join(", "),
                fk.references.table,
                fk.references.columns.join(", ")
            );
        }
    }
}

/// Count how many schema tables appear in the output text.
fn count_tables_in_output(output: &str, schema: &SchemaDefinition) -> usize {
    schema
        .tables
        .keys()
        .filter(|name| output.contains(name.as_str()))
        .count()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::schema::types::{ColumnDef, ForeignKeyDef, ForeignKeyRef, SqlDialect};
    use std::collections::HashMap;

    /// Helper to create a ColumnDef with defaults for synonyms.
    fn col(
        name: &str,
        data_type: &str,
        nullable: bool,
        desc: Option<&str>,
        tags: Vec<&str>,
    ) -> ColumnDef {
        ColumnDef {
            name: name.to_string(),
            data_type: data_type.to_string(),
            nullable,
            description: desc.map(|s| s.to_string()),
            tags: tags.into_iter().map(|s| s.to_string()).collect(),
            synonyms: vec![],

            statistics: None,
        }
    }

    fn make_test_schema() -> SchemaDefinition {
        let mut tables = HashMap::new();

        tables.insert(
            "users".to_string(),
            TableDef {
                description: Some("Core user table".to_string()),
                database: None,
                schema: None,
                columns: vec![
                    col("id", "INT", false, Some("User ID"), vec![]),
                    col("email", "VARCHAR", false, Some("User email"), vec!["pii"]),
                    col("name", "VARCHAR", true, Some("Full name"), vec!["pii"]),
                    col("created_at", "TIMESTAMP", false, None, vec![]),
                ],
                primary_key: Some(vec!["id".to_string()]),
                foreign_keys: vec![],
                statistics: None,

                clustering_keys: vec![],
            },
        );

        tables.insert(
            "orders".to_string(),
            TableDef {
                description: Some("Customer orders".to_string()),
                database: None,
                schema: None,
                columns: vec![
                    col("id", "INT", false, None, vec![]),
                    col("user_id", "INT", false, Some("FK to users.id"), vec![]),
                    col("total", "DECIMAL(10,2)", true, None, vec![]),
                    col("status", "VARCHAR", false, None, vec![]),
                ],
                primary_key: Some(vec!["id".to_string()]),
                foreign_keys: vec![ForeignKeyDef {
                    columns: vec!["user_id".to_string()],
                    references: ForeignKeyRef {
                        table: "users".to_string(),
                        columns: vec!["id".to_string()],
                    },
                }],
                statistics: None,

                clustering_keys: vec![],
            },
        );

        tables.insert(
            "products".to_string(),
            TableDef {
                description: Some("Product catalog".to_string()),
                database: None,
                schema: None,
                columns: vec![
                    col("id", "INT", false, None, vec![]),
                    col("name", "VARCHAR", false, None, vec![]),
                    col("price", "DECIMAL(10,2)", true, None, vec![]),
                ],
                primary_key: Some(vec!["id".to_string()]),
                foreign_keys: vec![],
                statistics: None,

                clustering_keys: vec![],
            },
        );

        tables.insert(
            "order_items".to_string(),
            TableDef {
                description: Some("Line items within orders".to_string()),
                database: None,
                schema: None,
                columns: vec![
                    col("id", "INT", false, None, vec![]),
                    col("order_id", "INT", false, None, vec![]),
                    col("product_id", "INT", false, None, vec![]),
                    col("quantity", "INT", false, None, vec![]),
                ],
                primary_key: Some(vec!["id".to_string()]),
                foreign_keys: vec![
                    ForeignKeyDef {
                        columns: vec!["order_id".to_string()],
                        references: ForeignKeyRef {
                            table: "orders".to_string(),
                            columns: vec!["id".to_string()],
                        },
                    },
                    ForeignKeyDef {
                        columns: vec!["product_id".to_string()],
                        references: ForeignKeyRef {
                            table: "products".to_string(),
                            columns: vec!["id".to_string()],
                        },
                    },
                ],
                statistics: None,

                clustering_keys: vec![],
            },
        );

        SchemaDefinition {
            version: "1".to_string(),
            dialect: SqlDialect::Generic,
            tables,
            database: None,
            schema_name: None,
            glossary: None,
        }
    }

    // ── Fingerprint tests ──

    #[test]
    fn test_fingerprint_is_deterministic() {
        let schema = make_test_schema();
        let fp1 = schema_fingerprint(&schema);
        let fp2 = schema_fingerprint(&schema);
        assert_eq!(fp1, fp2);
    }

    #[test]
    fn test_fingerprint_is_hex_string() {
        let schema = make_test_schema();
        let fp = schema_fingerprint(&schema);
        assert!(fp.chars().all(|c| c.is_ascii_hexdigit()));
        assert_eq!(fp.len(), 64); // SHA-256 = 64 hex chars
    }

    #[test]
    fn test_fingerprint_changes_on_schema_change() {
        let schema1 = make_test_schema();
        let mut schema2 = make_test_schema();
        schema2.tables.insert(
            "new_table".to_string(),
            TableDef {
                description: None,
                database: None,
                schema: None,
                columns: vec![],
                primary_key: None,
                foreign_keys: vec![],
                statistics: None,

                clustering_keys: vec![],
            },
        );
        assert_ne!(schema_fingerprint(&schema1), schema_fingerprint(&schema2));
    }

    // ── Token estimation tests ──

    #[test]
    fn test_estimate_tokens_empty() {
        assert_eq!(estimate_tokens(""), 0);
    }

    #[test]
    fn test_estimate_tokens_short() {
        // 4 chars => 1 token
        assert_eq!(estimate_tokens("abcd"), 1);
    }

    #[test]
    fn test_estimate_tokens_longer() {
        // 12 chars => 3 tokens
        assert_eq!(estimate_tokens("hello world!"), 3);
    }

    // ── Schema summary tests ──

    #[test]
    fn test_schema_summary_includes_all_tables() {
        let schema = make_test_schema();
        let summary = schema_summary(&schema);
        assert_eq!(summary.table_summaries.len(), 4);
    }

    #[test]
    fn test_schema_summary_natural_language_header() {
        let schema = make_test_schema();
        let summary = schema_summary(&schema);
        assert!(summary.natural_language.contains("4 tables"));
    }

    #[test]
    fn test_schema_summary_includes_relationships() {
        let schema = make_test_schema();
        let summary = schema_summary(&schema);
        assert!(summary.natural_language.contains("Key relationships"));
        assert!(
            summary
                .natural_language
                .contains("orders.user_id -> users.id")
        );
    }

    #[test]
    fn test_schema_summary_table_descriptions() {
        let schema = make_test_schema();
        let summary = schema_summary(&schema);
        let users_summary = summary
            .table_summaries
            .iter()
            .find(|t| t.name == "users")
            .expect("should have users");
        assert_eq!(users_summary.description, "Core user table");
        assert_eq!(users_summary.column_count, 4);
    }

    #[test]
    fn test_schema_summary_key_columns() {
        let schema = make_test_schema();
        let summary = schema_summary(&schema);
        let orders_summary = summary
            .table_summaries
            .iter()
            .find(|t| t.name == "orders")
            .expect("should have orders");
        assert!(orders_summary.key_columns.contains(&"id".to_string()));
        assert!(orders_summary.key_columns.contains(&"user_id".to_string()));
    }

    #[test]
    fn test_schema_summary_estimated_tokens() {
        let schema = make_test_schema();
        let summary = schema_summary(&schema);
        assert!(summary.estimated_tokens > 0);
    }

    // ── Domain inference tests ──

    #[test]
    fn test_infer_domain_ecommerce() {
        let schema = make_test_schema();
        let domain = infer_domain(&schema);
        assert_eq!(domain, "e-commerce");
    }

    #[test]
    fn test_infer_domain_empty() {
        let schema = SchemaDefinition {
            version: "1".to_string(),
            dialect: SqlDialect::Generic,
            tables: HashMap::new(),
            database: None,
            schema_name: None,
            glossary: None,
        };
        let domain = infer_domain(&schema);
        assert!(domain.is_empty());
    }

    // ── Disclosure level tests ──

    #[test]
    fn test_disclosure_level_ordering() {
        assert!(DisclosureLevel::Fingerprint < DisclosureLevel::TableList);
        assert!(DisclosureLevel::TableList < DisclosureLevel::ColumnNames);
        assert!(DisclosureLevel::ColumnNames < DisclosureLevel::FullRelevant);
        assert!(DisclosureLevel::FullRelevant < DisclosureLevel::Full);
    }

    #[test]
    fn test_disclosure_level_next() {
        assert_eq!(
            DisclosureLevel::Fingerprint.next(),
            Some(DisclosureLevel::TableList)
        );
        assert_eq!(DisclosureLevel::Full.next(), None);
    }

    #[test]
    fn test_disclosure_level_prev() {
        assert_eq!(DisclosureLevel::Fingerprint.prev(), None);
        assert_eq!(
            DisclosureLevel::Full.prev(),
            Some(DisclosureLevel::FullRelevant)
        );
    }

    #[test]
    fn test_disclosure_level_as_u8() {
        assert_eq!(DisclosureLevel::Fingerprint.as_u8(), 0);
        assert_eq!(DisclosureLevel::Full.as_u8(), 4);
    }

    // ── optimize_context tests ──

    #[test]
    fn test_optimize_context_fingerprint_level() {
        let schema = make_test_schema();
        let config = ContextConfig {
            level: DisclosureLevel::Fingerprint,
            ..Default::default()
        };
        let result = optimize_context(&schema, &config).expect("should optimize");
        assert_eq!(result.level_used, DisclosureLevel::Fingerprint);
        assert!(result.compressed_schema.contains("schema_fingerprint"));
        assert_eq!(result.tables_total, 4);
    }

    #[test]
    fn test_optimize_context_table_list_level() {
        let schema = make_test_schema();
        let config = ContextConfig {
            level: DisclosureLevel::TableList,
            max_tokens: 5000,
            ..Default::default()
        };
        let result = optimize_context(&schema, &config).expect("should optimize");
        assert_eq!(result.level_used, DisclosureLevel::TableList);
        assert!(result.compressed_schema.contains("users"));
        assert!(result.compressed_schema.contains("orders"));
    }

    #[test]
    fn test_optimize_context_column_names_level() {
        let schema = make_test_schema();
        let config = ContextConfig {
            level: DisclosureLevel::ColumnNames,
            max_tokens: 5000,
            ..Default::default()
        };
        let result = optimize_context(&schema, &config).expect("should optimize");
        assert_eq!(result.level_used, DisclosureLevel::ColumnNames);
        assert!(result.compressed_schema.contains("email"));
    }

    #[test]
    fn test_optimize_context_full_level() {
        let schema = make_test_schema();
        let config = ContextConfig {
            level: DisclosureLevel::Full,
            max_tokens: 50000,
            ..Default::default()
        };
        let result = optimize_context(&schema, &config).expect("should optimize");
        assert_eq!(result.level_used, DisclosureLevel::Full);
        assert!(result.compressed_schema.contains("INT"));
        assert!(result.compressed_schema.contains("VARCHAR"));
    }

    #[test]
    fn test_optimize_context_auto_stepdown() {
        let schema = make_test_schema();
        // Very small budget should force stepdown
        let config = ContextConfig {
            level: DisclosureLevel::Full,
            max_tokens: 10,
            ..Default::default()
        };
        let result = optimize_context(&schema, &config).expect("should optimize");
        // Should have stepped down to a lower level
        assert!(result.level_used < DisclosureLevel::Full);
    }

    #[test]
    fn test_optimize_context_compression_ratio() {
        let schema = make_test_schema();
        let config = ContextConfig {
            level: DisclosureLevel::TableList,
            max_tokens: 5000,
            ..Default::default()
        };
        let result = optimize_context(&schema, &config).expect("should optimize");
        assert!(result.compression_ratio > 0.0);
        assert!(result.compression_ratio < 1.0);
    }

    #[test]
    fn test_optimize_context_fingerprint_included() {
        let schema = make_test_schema();
        let config = ContextConfig::default();
        let result = optimize_context(&schema, &config).expect("should optimize");
        assert!(!result.fingerprint.is_empty());
        assert_eq!(result.fingerprint.len(), 64);
    }

    // ── optimize_for_query tests ──

    #[test]
    fn test_optimize_for_query_single_table() {
        let schema = make_test_schema();
        let config = ContextConfig {
            level: DisclosureLevel::FullRelevant,
            max_tokens: 5000,
            ..Default::default()
        };
        let result =
            optimize_for_query("SELECT id FROM users", &schema, &config).expect("should optimize");
        // Should include users at full detail
        assert!(result.compressed_schema.contains("users"));
        assert!(result.tables_included <= result.tables_total);
    }

    #[test]
    fn test_optimize_for_query_with_join() {
        let schema = make_test_schema();
        let config = ContextConfig {
            level: DisclosureLevel::FullRelevant,
            max_tokens: 5000,
            ..Default::default()
        };
        let result = optimize_for_query(
            "SELECT u.name, o.total FROM users u JOIN orders o ON u.id = o.user_id",
            &schema,
            &config,
        )
        .expect("should optimize");
        assert!(result.compressed_schema.contains("users"));
        assert!(result.compressed_schema.contains("orders"));
    }

    #[test]
    fn test_optimize_for_query_fk_expansion() {
        let schema = make_test_schema();
        let config = ContextConfig {
            level: DisclosureLevel::FullRelevant,
            max_tokens: 5000,
            ..Default::default()
        };
        // Querying order_items should pull in orders and products via FK
        let result = optimize_for_query("SELECT quantity FROM order_items", &schema, &config)
            .expect("should optimize");
        assert!(result.compressed_schema.contains("order_items"));
    }

    #[test]
    fn test_optimize_for_query_priority_tables() {
        let schema = make_test_schema();
        let config = ContextConfig {
            level: DisclosureLevel::FullRelevant,
            max_tokens: 5000,
            priority_tables: vec!["products".to_string()],
            ..Default::default()
        };
        let result =
            optimize_for_query("SELECT id FROM users", &schema, &config).expect("should optimize");
        // Products should be included even though it's not in the query
        assert!(result.compressed_schema.contains("products"));
    }

    #[test]
    fn test_optimize_for_query_no_tables() {
        let schema = make_test_schema();
        let config = ContextConfig {
            level: DisclosureLevel::FullRelevant,
            max_tokens: 5000,
            ..Default::default()
        };
        let result = optimize_for_query("SELECT 1 + 1", &schema, &config).expect("should optimize");
        // No relevant tables, but all tables should still appear in compact form
        assert!(result.tables_total == 4);
    }

    // ── Config tests ──

    #[test]
    fn test_default_config() {
        let config = ContextConfig::default();
        assert_eq!(config.max_tokens, 2000);
        assert_eq!(config.level, DisclosureLevel::ColumnNames);
        assert!(config.include_descriptions);
        assert!(!config.include_tags);
        assert!(config.include_foreign_keys);
        assert!(config.priority_tables.is_empty());
    }

    #[test]
    fn test_config_without_descriptions() {
        let schema = make_test_schema();
        let config = ContextConfig {
            level: DisclosureLevel::Full,
            max_tokens: 50000,
            include_descriptions: false,
            ..Default::default()
        };
        let result = optimize_context(&schema, &config).expect("should optimize");
        // "Core user table" description should not appear
        assert!(!result.compressed_schema.contains("Core user table"));
    }

    #[test]
    fn test_config_with_tags() {
        let schema = make_test_schema();
        let config = ContextConfig {
            level: DisclosureLevel::Full,
            max_tokens: 50000,
            include_tags: true,
            ..Default::default()
        };
        let result = optimize_context(&schema, &config).expect("should optimize");
        assert!(result.compressed_schema.contains("[pii]"));
    }

    #[test]
    fn test_config_without_foreign_keys() {
        let schema = make_test_schema();
        let config = ContextConfig {
            level: DisclosureLevel::Full,
            max_tokens: 50000,
            include_foreign_keys: false,
            ..Default::default()
        };
        let result = optimize_context(&schema, &config).expect("should optimize");
        assert!(!result.compressed_schema.contains("foreign_keys:"));
    }

    // ── Rendering tests ──

    #[test]
    fn test_render_fingerprint_compact() {
        let schema = make_test_schema();
        let result = render_at_level(
            &schema,
            DisclosureLevel::Fingerprint,
            &ContextConfig::default(),
            None,
        )
        .expect("should render");
        assert!(result.contains("schema_fingerprint:"));
        assert!(result.contains("tables: 4"));
    }

    #[test]
    fn test_render_table_list_format() {
        let schema = make_test_schema();
        let result = render_at_level(
            &schema,
            DisclosureLevel::TableList,
            &ContextConfig::default(),
            None,
        )
        .expect("should render");
        assert!(result.contains("4 cols"));
        assert!(result.contains("Core user table"));
    }

    #[test]
    fn test_render_column_names_format() {
        let schema = make_test_schema();
        let result = render_at_level(
            &schema,
            DisclosureLevel::ColumnNames,
            &ContextConfig::default(),
            None,
        )
        .expect("should render");
        assert!(result.contains("columns: ["));
        assert!(result.contains("email"));
    }

    #[test]
    fn test_render_full_includes_types() {
        let schema = make_test_schema();
        let result = render_at_level(
            &schema,
            DisclosureLevel::Full,
            &ContextConfig::default(),
            None,
        )
        .expect("should render");
        assert!(result.contains("INT"));
        assert!(result.contains("VARCHAR"));
        assert!(result.contains("DECIMAL(10,2)"));
    }

    #[test]
    fn test_render_full_relevant_mixed_detail() {
        let schema = make_test_schema();
        let mut relevant = HashSet::new();
        relevant.insert("users".to_string());
        let result = render_at_level(
            &schema,
            DisclosureLevel::FullRelevant,
            &ContextConfig::default(),
            Some(&relevant),
        )
        .expect("should render");
        // users should have full detail with types
        assert!(result.contains("- id: INT"));
        // orders should be compact
        assert!(result.contains("orders: ["));
    }

    // ── Empty schema tests ──

    #[test]
    fn test_optimize_empty_schema() {
        let schema = SchemaDefinition {
            version: "1".to_string(),
            dialect: SqlDialect::Generic,
            tables: HashMap::new(),
            database: None,
            schema_name: None,
            glossary: None,
        };
        let config = ContextConfig::default();
        let result = optimize_context(&schema, &config).expect("should optimize");
        assert_eq!(result.tables_total, 0);
        assert_eq!(result.tables_included, 0);
    }

    #[test]
    fn test_summary_empty_schema() {
        let schema = SchemaDefinition {
            version: "1".to_string(),
            dialect: SqlDialect::Generic,
            tables: HashMap::new(),
            database: None,
            schema_name: None,
            glossary: None,
        };
        let summary = schema_summary(&schema);
        assert!(summary.natural_language.contains("0 tables"));
        assert!(summary.table_summaries.is_empty());
    }
}
