//! Phase-85: Natural Language to SQL Engine
//!
//! Combines the LLM Gateway with the DataFusion SQL engine to convert
//! natural-language questions into SQL, execute them, and return human-readable
//! answers.
//!
//! # Architecture
//!
//! ```text
//! User question (NL)
//!     │
//!     ▼
//! NlToSqlEngine::ask()
//!     ├─ 1. Collect schema context (tables + columns + types)
//!     ├─ 2. Build prompt (schema + question → LLM)
//!     ├─ 3. LLM generates SQL
//!     ├─ 4. SQL safety validation (whitelist, depth, row limit)
//!     ├─ 5. Execute SQL via DataFusion
//!     ├─ 6. On error → multi-round correction (up to max_retries)
//!     └─ 7. Format result → optional NL answer via LLM
//! ```

use std::fmt;

use serde::{Deserialize, Serialize};

use super::client::{LlmClient, LlmMessage};
use super::error::LlmResult;
use super::types::ChatConfig;

// ============================================================================
// Configuration
// ============================================================================

/// Configuration for the NL-to-SQL engine.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NlToSqlConfig {
    /// Maximum retry attempts when SQL execution fails (default: 3).
    #[serde(default = "default_max_retries")]
    pub max_retries: usize,

    /// Maximum rows to return from a query (default: 1000).
    #[serde(default = "default_max_rows")]
    pub max_rows: usize,

    /// Whether to generate a natural-language answer from results (default: true).
    #[serde(default = "default_true")]
    pub generate_nl_answer: bool,

    /// Maximum SQL sub-query nesting depth allowed (default: 3).
    #[serde(default = "default_max_depth")]
    pub max_subquery_depth: usize,

    /// Temperature for LLM SQL generation (default: 0.0 for deterministic).
    #[serde(default)]
    pub temperature: f32,

    /// Model name override (uses router default if None).
    pub model: Option<String>,
}

fn default_max_retries() -> usize { 3 }
fn default_max_rows() -> usize { 1000 }
fn default_true() -> bool { true }
fn default_max_depth() -> usize { 3 }

impl Default for NlToSqlConfig {
    fn default() -> Self {
        Self {
            max_retries: default_max_retries(),
            max_rows: default_max_rows(),
            generate_nl_answer: true,
            max_subquery_depth: default_max_depth(),
            temperature: 0.0f32,
            model: None,
        }
    }
}

// ============================================================================
// SQL Safety Validation
// ============================================================================

/// Result of SQL safety validation.
#[derive(Debug, Clone)]
pub enum SqlSafetyResult {
    /// SQL is safe to execute.
    Safe,
    /// SQL is rejected with a reason.
    Rejected(String),
}

impl SqlSafetyResult {
    pub fn is_safe(&self) -> bool {
        matches!(self, SqlSafetyResult::Safe)
    }
}

/// SQL statement categories that are allowed.
const ALLOWED_PREFIXES: &[&str] = &[
    "SELECT", "SHOW", "DESCRIBE", "EXPLAIN", "WITH",
];

/// Dangerous keywords that are blocked.
const BLOCKED_KEYWORDS: &[&str] = &[
    "DROP", "DELETE", "UPDATE", "INSERT", "ALTER", "TRUNCATE",
    "CREATE", "GRANT", "REVOKE", "EXEC", "EXECUTE",
    "INTO OUTFILE", "LOAD DATA", "LOAD_FILE",
];

/// Validate SQL safety.
///
/// Rules:
/// 1. Must start with an allowed prefix (SELECT, SHOW, DESCRIBE, EXPLAIN, WITH)
/// 2. Must NOT contain blocked keywords as standalone tokens
/// 3. Sub-query nesting depth must not exceed `max_depth`
/// 4. Must contain at least one meaningful token
pub fn validate_sql_safety(sql: &str, max_depth: usize) -> SqlSafetyResult {
    let trimmed = sql.trim();
    if trimmed.is_empty() {
        return SqlSafetyResult::Rejected("Empty SQL statement".to_string());
    }

    // Strip markdown code fences if present
    let clean = strip_sql_fences(trimmed);
    let upper = clean.to_uppercase();

    // 1. Check allowed prefix
    let has_allowed_prefix = ALLOWED_PREFIXES.iter().any(|p| upper.starts_with(p));
    if !has_allowed_prefix {
        return SqlSafetyResult::Rejected(format!(
            "SQL must start with one of: {}. Got: '{}'",
            ALLOWED_PREFIXES.join(", "),
            &upper[..upper.len().min(30)]
        ));
    }

    // 2. Check blocked keywords (as whole words, not substrings)
    for keyword in BLOCKED_KEYWORDS {
        // Use word boundary check
        if contains_keyword(&upper, keyword) {
            return SqlSafetyResult::Rejected(format!(
                "SQL contains blocked keyword: {}",
                keyword
            ));
        }
    }

    // 3. Check sub-query nesting depth
    let depth = count_subquery_depth(&upper);
    if depth > max_depth {
        return SqlSafetyResult::Rejected(format!(
            "Sub-query nesting depth {} exceeds maximum {}",
            depth, max_depth
        ));
    }

    SqlSafetyResult::Safe
}

/// Strip markdown SQL code fences (```sql ... ```)
fn strip_sql_fences(sql: &str) -> &str {
    let s = sql.trim();
    if s.starts_with("```") {
        // Find end of first line
        let after_fence = if let Some(nl) = s.find('\n') {
            &s[nl + 1..]
        } else {
            return s;
        };
        // Strip trailing fence
        let content = if after_fence.trim_end().ends_with("```") {
            let end = after_fence.rfind("```").unwrap_or(after_fence.len());
            &after_fence[..end]
        } else {
            after_fence
        };
        content.trim()
    } else {
        s
    }
}

/// Check if a keyword appears as a standalone word in the SQL.
fn contains_keyword(upper_sql: &str, keyword: &str) -> bool {
    let mut search_from = 0;
    while let Some(pos) = upper_sql[search_from..].find(keyword) {
        let abs_pos = search_from + pos;
        let before_ok = abs_pos == 0
            || !upper_sql.as_bytes()[abs_pos - 1].is_ascii_alphanumeric();
        let after_pos = abs_pos + keyword.len();
        let after_ok = after_pos >= upper_sql.len()
            || !upper_sql.as_bytes()[after_pos].is_ascii_alphanumeric();

        if before_ok && after_ok {
            // But skip if inside a string literal
            if !is_inside_string_literal(upper_sql, abs_pos) {
                return true;
            }
        }
        search_from = abs_pos + keyword.len();
    }
    false
}

/// Simple heuristic: count single quotes before position; if odd, we're inside a string.
fn is_inside_string_literal(sql: &str, pos: usize) -> bool {
    let before = &sql[..pos];
    before.chars().filter(|&c| c == '\'').count() % 2 != 0
}

/// Count maximum sub-query nesting depth by tracking SELECT ... ( SELECT ... )
fn count_subquery_depth(upper_sql: &str) -> usize {
    let mut max_depth = 0usize;
    let mut current_depth = 0usize;

    // Simple approach: count depth by parentheses containing SELECT
    let mut paren_depth = 0usize;
    let chars: Vec<char> = upper_sql.chars().collect();
    let len = chars.len();
    let mut i = 0;

    while i < len {
        match chars[i] {
            '(' => {
                paren_depth += 1;
                // Check if SELECT follows after optional whitespace
                let rest: String = chars[i + 1..].iter().collect();
                let trimmed = rest.trim_start();
                if trimmed.starts_with("SELECT") || trimmed.starts_with("WITH") {
                    current_depth += 1;
                    if current_depth > max_depth {
                        max_depth = current_depth;
                    }
                }
            }
            ')' => {
                if paren_depth > 0 {
                    paren_depth -= 1;
                    if current_depth > 0 {
                        current_depth -= 1;
                    }
                }
            }
            _ => {}
        }
        i += 1;
    }

    max_depth
}

// ============================================================================
// Schema Context
// ============================================================================

/// Description of a table column for schema context injection.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ColumnInfo {
    pub name: String,
    pub data_type: String,
    pub nullable: bool,
}

impl fmt::Display for ColumnInfo {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let null_str = if self.nullable { "NULL" } else { "NOT NULL" };
        write!(f, "  {} {} {}", self.name, self.data_type, null_str)
    }
}

/// Description of a table for schema context injection.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TableInfo {
    pub name: String,
    pub columns: Vec<ColumnInfo>,
    pub row_count: Option<usize>,
}

impl fmt::Display for TableInfo {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let rows = self.row_count.map(|r| format!(" (~{} rows)", r)).unwrap_or_default();
        writeln!(f, "TABLE: {}{}", self.name, rows)?;
        for col in &self.columns {
            writeln!(f, "{}", col)?;
        }
        Ok(())
    }
}

/// Schema context for all registered tables.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct SchemaContext {
    pub tables: Vec<TableInfo>,
}

impl SchemaContext {
    pub fn new() -> Self {
        Self { tables: Vec::new() }
    }

    pub fn add_table(&mut self, table: TableInfo) {
        self.tables.push(table);
    }

    pub fn is_empty(&self) -> bool {
        self.tables.is_empty()
    }

    /// Format as a concise schema description for LLM prompt injection.
    pub fn to_prompt_context(&self) -> String {
        if self.tables.is_empty() {
            return "No tables are registered.".to_string();
        }

        let mut out = String::from("## Available Database Tables\n\n");
        for table in &self.tables {
            out.push_str(&format!("{}\n", table));
        }
        out
    }
}

// ============================================================================
// NL-to-SQL Query Result
// ============================================================================

/// Result of a natural-language query.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NlQueryResult {
    /// The original user question.
    pub question: String,

    /// The generated SQL statement.
    pub sql: String,

    /// Whether the SQL was executed.
    pub executed: bool,

    /// Query result as JSON rows (if executed).
    #[serde(default)]
    pub data: Vec<serde_json::Value>,

    /// Total number of result rows.
    pub row_count: usize,

    /// Natural language answer (if generated).
    pub answer: Option<String>,

    /// SQL explanation (if generated).
    pub explanation: Option<String>,

    /// Execution time in milliseconds.
    pub execution_time_ms: f64,

    /// Number of correction rounds used (0 = first attempt succeeded).
    pub correction_rounds: usize,

    /// History of SQL attempts (including failed ones).
    #[serde(default)]
    pub sql_history: Vec<SqlAttempt>,
}

/// Record of a single SQL generation/execution attempt.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SqlAttempt {
    pub sql: String,
    pub success: bool,
    pub error: Option<String>,
}

// ============================================================================
// NlToSqlEngine
// ============================================================================

/// Natural Language to SQL query engine.
///
/// Bridges the LLM Gateway with the DataFusion SQL engine to:
/// 1. Collect schema context from registered tables
/// 2. Generate SQL from natural-language questions via LLM
/// 3. Validate SQL safety
/// 4. Execute SQL and handle errors with multi-round correction
/// 5. Optionally generate natural-language answers from results
pub struct NlToSqlEngine {
    /// Configuration.
    config: NlToSqlConfig,
    /// Schema context (cached).
    schema: SchemaContext,
    /// Conversation history for multi-turn support.
    history: Vec<(String, String)>, // (question, sql)
}

impl NlToSqlEngine {
    /// Create a new NL-to-SQL engine with default configuration.
    pub fn new() -> Self {
        Self {
            config: NlToSqlConfig::default(),
            schema: SchemaContext::new(),
            history: Vec::new(),
        }
    }

    /// Create with custom configuration.
    pub fn with_config(config: NlToSqlConfig) -> Self {
        Self {
            config,
            schema: SchemaContext::new(),
            history: Vec::new(),
        }
    }

    /// Get the current configuration.
    pub fn config(&self) -> &NlToSqlConfig {
        &self.config
    }

    /// Set schema context (tables + columns).
    pub fn set_schema(&mut self, schema: SchemaContext) {
        self.schema = schema;
    }

    /// Get the current schema context.
    pub fn schema(&self) -> &SchemaContext {
        &self.schema
    }

    /// Add a table to the schema context.
    pub fn add_table(&mut self, table: TableInfo) {
        self.schema.add_table(table);
    }

    /// Clear conversation history.
    pub fn clear_history(&mut self) {
        self.history.clear();
    }

    /// Get conversation history.
    pub fn history(&self) -> &[(String, String)] {
        &self.history
    }

    // ── Core: Generate SQL ──────────────────────────────────────────────

    /// Generate SQL from a natural-language question using an LLM.
    ///
    /// Returns the generated SQL string (not yet executed).
    pub async fn generate_sql(
        &self,
        question: &str,
        client: &dyn LlmClient,
    ) -> LlmResult<String> {
        let messages = self.build_generate_messages(question, None);

        let config = ChatConfig {
            temperature: Some(self.config.temperature),
            max_tokens: Some(2048),
            ..Default::default()
        };

        let response = client.chat(&messages, &config).await?;
        let sql = extract_sql_from_response(&response.content);
        Ok(sql)
    }

    /// Generate SQL with an error feedback message for correction.
    pub async fn generate_sql_with_correction(
        &self,
        question: &str,
        previous_sql: &str,
        error_msg: &str,
        client: &dyn LlmClient,
    ) -> LlmResult<String> {
        let correction = format!(
            "The previous SQL was:\n```sql\n{}\n```\n\nIt failed with error:\n{}\n\nPlease fix the SQL and try again.",
            previous_sql, error_msg
        );
        let messages = self.build_generate_messages(question, Some(&correction));

        let config = ChatConfig {
            temperature: Some(self.config.temperature),
            max_tokens: Some(2048),
            ..Default::default()
        };

        let response = client.chat(&messages, &config).await?;
        let sql = extract_sql_from_response(&response.content);
        Ok(sql)
    }

    // ── Core: Full Pipeline (generate + validate + execute) ─────────────

    /// Execute the full NL-to-SQL pipeline:
    /// 1. Generate SQL from question
    /// 2. Validate safety
    /// 3. Execute via the provided executor function
    /// 4. On failure, retry with error feedback (up to max_retries)
    /// 5. Optionally generate NL answer
    ///
    /// The `execute_fn` is called with the SQL string and should return
    /// either the rows as JSON or an error string.
    pub async fn ask<F>(
        &mut self,
        question: &str,
        client: &dyn LlmClient,
        execute_fn: F,
    ) -> LlmResult<NlQueryResult>
    where
        F: Fn(&str) -> Result<(Vec<serde_json::Value>, usize), String>,
    {
        let start = std::time::Instant::now();
        let mut sql_history = Vec::new();
        let mut correction_rounds = 0;

        // 1. Generate initial SQL
        let mut sql = self.generate_sql(question, client).await?;

        loop {
            // 2. Validate safety
            let safety = validate_sql_safety(&sql, self.config.max_subquery_depth);
            if let SqlSafetyResult::Rejected(reason) = safety {
                sql_history.push(SqlAttempt {
                    sql: sql.clone(),
                    success: false,
                    error: Some(reason.clone()),
                });

                if correction_rounds >= self.config.max_retries {
                    return Ok(NlQueryResult {
                        question: question.to_string(),
                        sql,
                        executed: false,
                        data: Vec::new(),
                        row_count: 0,
                        answer: Some(format!("SQL safety validation failed: {}", reason)),
                        explanation: None,
                        execution_time_ms: start.elapsed().as_secs_f64() * 1000.0,
                        correction_rounds,
                        sql_history,
                    });
                }

                correction_rounds += 1;
                sql = self
                    .generate_sql_with_correction(question, &sql, &reason, client)
                    .await?;
                continue;
            }

            // 3. Apply LIMIT if not already present
            let exec_sql = ensure_limit(&sql, self.config.max_rows);

            // 4. Execute SQL
            match execute_fn(&exec_sql) {
                Ok((data, row_count)) => {
                    sql_history.push(SqlAttempt {
                        sql: exec_sql.clone(),
                        success: true,
                        error: None,
                    });

                    // 5. Optionally generate NL answer
                    let answer = if self.config.generate_nl_answer && !data.is_empty() {
                        match self
                            .generate_nl_answer(question, &exec_sql, &data, client)
                            .await
                        {
                            Ok(ans) => Some(ans),
                            Err(_) => None,
                        }
                    } else if data.is_empty() {
                        Some("No results found for your query.".to_string())
                    } else {
                        None
                    };

                    // Record in history for multi-turn
                    self.history.push((question.to_string(), exec_sql.clone()));

                    return Ok(NlQueryResult {
                        question: question.to_string(),
                        sql: exec_sql,
                        executed: true,
                        data,
                        row_count,
                        answer,
                        explanation: None,
                        execution_time_ms: start.elapsed().as_secs_f64() * 1000.0,
                        correction_rounds,
                        sql_history,
                    });
                }
                Err(error_msg) => {
                    sql_history.push(SqlAttempt {
                        sql: exec_sql.clone(),
                        success: false,
                        error: Some(error_msg.clone()),
                    });

                    if correction_rounds >= self.config.max_retries {
                        return Ok(NlQueryResult {
                            question: question.to_string(),
                            sql: exec_sql,
                            executed: false,
                            data: Vec::new(),
                            row_count: 0,
                            answer: Some(format!(
                                "Failed to execute SQL after {} attempts. Last error: {}",
                                correction_rounds + 1,
                                error_msg
                            )),
                            explanation: None,
                            execution_time_ms: start.elapsed().as_secs_f64() * 1000.0,
                            correction_rounds,
                            sql_history,
                        });
                    }

                    correction_rounds += 1;
                    sql = self
                        .generate_sql_with_correction(question, &exec_sql, &error_msg, client)
                        .await?;
                }
            }
        }
    }

    /// Preview mode: only generate SQL without executing.
    pub async fn preview(
        &self,
        question: &str,
        client: &dyn LlmClient,
    ) -> LlmResult<NlQueryResult> {
        let start = std::time::Instant::now();
        let sql = self.generate_sql(question, client).await?;

        let safety = validate_sql_safety(&sql, self.config.max_subquery_depth);
        let safe = safety.is_safe();
        let safety_msg = match &safety {
            SqlSafetyResult::Safe => None,
            SqlSafetyResult::Rejected(reason) => Some(reason.clone()),
        };

        Ok(NlQueryResult {
            question: question.to_string(),
            sql,
            executed: false,
            data: Vec::new(),
            row_count: 0,
            answer: safety_msg,
            explanation: if safe {
                Some("SQL generated and validated. Ready to execute.".to_string())
            } else {
                Some("SQL failed safety validation.".to_string())
            },
            execution_time_ms: start.elapsed().as_secs_f64() * 1000.0,
            correction_rounds: 0,
            sql_history: Vec::new(),
        })
    }

    // ── Internal: Prompt Building ───────────────────────────────────────

    fn build_generate_messages(
        &self,
        question: &str,
        correction: Option<&str>,
    ) -> Vec<LlmMessage> {
        let mut messages = Vec::new();

        // System prompt with schema context
        let system_prompt = self.build_system_prompt();
        messages.push(LlmMessage::system(system_prompt));

        // Add conversation history for multi-turn context
        for (prev_q, prev_sql) in &self.history {
            messages.push(LlmMessage::user(prev_q.clone()));
            messages.push(LlmMessage::assistant(format!(
                "```sql\n{}\n```",
                prev_sql
            )));
        }

        // Current question
        messages.push(LlmMessage::user(question.to_string()));

        // Correction feedback (if retrying)
        if let Some(corr) = correction {
            messages.push(LlmMessage::assistant("I'll fix the SQL.".to_string()));
            messages.push(LlmMessage::user(corr.to_string()));
        }

        messages
    }

    fn build_system_prompt(&self) -> String {
        let schema_ctx = self.schema.to_prompt_context();

        format!(
            r#"You are a SQL expert assistant. Your job is to convert natural language questions into valid SQL queries.

## Rules
1. Generate ONLY valid SQL. Output ONLY the SQL statement, wrapped in ```sql ... ``` code fences.
2. Only use SELECT, SHOW, DESCRIBE, or EXPLAIN statements. Never use INSERT, UPDATE, DELETE, DROP, or any DDL/DML.
3. Always include a LIMIT clause (max 1000 rows) unless the user explicitly asks for all data.
4. Use table and column names exactly as shown in the schema below.
5. For text search, use LIKE with % wildcards.
6. For aggregations, use GROUP BY and appropriate aggregate functions.
7. If the question is ambiguous, make a reasonable assumption and note it.

## Database Schema

{schema_ctx}

## Output Format
Respond with ONLY the SQL query in a code block. No explanations before or after."#
        )
    }

    /// Generate a natural-language answer from SQL results.
    async fn generate_nl_answer(
        &self,
        question: &str,
        sql: &str,
        data: &[serde_json::Value],
        client: &dyn LlmClient,
    ) -> LlmResult<String> {
        // Truncate data to avoid token overflow
        let preview_data: Vec<_> = data.iter().take(20).collect();
        let data_str = serde_json::to_string_pretty(&preview_data).unwrap_or_default();
        let total = data.len();

        let prompt = format!(
            r#"Based on the following SQL query result, provide a concise natural-language answer to the user's question.

Question: {question}

SQL: {sql}

Result ({total} rows, showing first {shown}):
{data_str}

Provide a clear, concise answer in 1-3 sentences. If the data contains numbers, include them in your answer."#,
            shown = preview_data.len()
        );

        let messages = vec![
            LlmMessage::system(
                "You are a helpful data analyst. Summarize SQL query results into clear, concise answers."
                    .to_string(),
            ),
            LlmMessage::user(prompt),
        ];

        let config = ChatConfig {
            temperature: Some(0.3f32),
            max_tokens: Some(512),
            ..Default::default()
        };

        let response = client.chat(&messages, &config).await?;
        Ok(response.content)
    }
}

impl Default for NlToSqlEngine {
    fn default() -> Self {
        Self::new()
    }
}

// ============================================================================
// Helpers
// ============================================================================

/// Extract SQL from LLM response (strip markdown fences and explanations).
pub fn extract_sql_from_response(response: &str) -> String {
    let trimmed = response.trim();

    // Try to extract from ```sql ... ``` block
    if let Some(start) = trimmed.find("```sql") {
        let after = &trimmed[start + 6..];
        if let Some(end) = after.find("```") {
            return after[..end].trim().to_string();
        }
    }

    // Try to extract from ``` ... ``` block (without language tag)
    if let Some(start) = trimmed.find("```") {
        let after = &trimmed[start + 3..];
        // Skip optional language tag on the first line
        let content_start = if let Some(nl) = after.find('\n') {
            nl + 1
        } else {
            0
        };
        let content = &after[content_start..];
        if let Some(end) = content.find("```") {
            return content[..end].trim().to_string();
        }
    }

    // No code fence found — try to extract the first line that looks like SQL
    for line in trimmed.lines() {
        let l = line.trim().to_uppercase();
        if l.starts_with("SELECT") || l.starts_with("WITH") || l.starts_with("SHOW") || l.starts_with("DESCRIBE") || l.starts_with("EXPLAIN") {
            // Return from this line to end (multi-line SQL)
            if let Some(pos) = trimmed.find(line) {
                return trimmed[pos..].trim().to_string();
            }
        }
    }

    // Fallback: return as-is
    trimmed.to_string()
}

/// Ensure the SQL has a LIMIT clause. If not, append one.
pub fn ensure_limit(sql: &str, max_rows: usize) -> String {
    let upper = sql.trim().to_uppercase();

    // Don't add LIMIT to SHOW/DESCRIBE/EXPLAIN
    if upper.starts_with("SHOW") || upper.starts_with("DESCRIBE") || upper.starts_with("EXPLAIN") {
        return sql.to_string();
    }

    // Check if LIMIT already exists (simple heuristic)
    if upper.contains(" LIMIT ") {
        // Parse existing limit and cap it
        return cap_existing_limit(sql, max_rows);
    }

    // Remove trailing semicolons
    let clean = sql.trim().trim_end_matches(';').trim();
    format!("{} LIMIT {}", clean, max_rows)
}

/// Cap an existing LIMIT value if it exceeds max_rows.
fn cap_existing_limit(sql: &str, max_rows: usize) -> String {
    let upper = sql.to_uppercase();
    if let Some(limit_pos) = upper.rfind(" LIMIT ") {
        let after_limit = &sql[limit_pos + 7..];
        let num_str: String = after_limit.trim().chars().take_while(|c| c.is_ascii_digit()).collect();
        if let Ok(existing) = num_str.parse::<usize>() {
            if existing > max_rows {
                let before = &sql[..limit_pos];
                let after_num = &after_limit[num_str.len()..];
                return format!("{} LIMIT {}{}", before.trim(), max_rows, after_num);
            }
        }
    }
    sql.to_string()
}

/// Convert Arrow RecordBatches to JSON rows.
#[cfg(feature = "sql-engine")]
pub fn record_batches_to_json(
    batches: &[arrow::array::RecordBatch],
) -> Result<Vec<serde_json::Value>, String> {
    use arrow::json::writer::record_batches_to_json_rows;

    let mut all_rows = Vec::new();
    for batch in batches {
        let rows = record_batches_to_json_rows(&[batch])
            .map_err(|e| format!("JSON conversion error: {}", e))?;
        for row in rows {
            all_rows.push(serde_json::Value::Object(row));
        }
    }
    Ok(all_rows)
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    // ── Config tests ─────────────────────────────────────────────────────

    #[test]
    fn test_config_defaults() {
        let cfg = NlToSqlConfig::default();
        assert_eq!(cfg.max_retries, 3);
        assert_eq!(cfg.max_rows, 1000);
        assert!(cfg.generate_nl_answer);
        assert_eq!(cfg.max_subquery_depth, 3);
        assert_eq!(cfg.temperature, 0.0f32);
        assert!(cfg.model.is_none());
    }

    // ── SQL safety validation tests ──────────────────────────────────────

    #[test]
    fn test_safe_select() {
        assert!(validate_sql_safety("SELECT * FROM users", 3).is_safe());
    }

    #[test]
    fn test_safe_with_cte() {
        assert!(validate_sql_safety(
            "WITH cte AS (SELECT id FROM users) SELECT * FROM cte",
            3
        ).is_safe());
    }

    #[test]
    fn test_safe_show() {
        assert!(validate_sql_safety("SHOW TABLES", 3).is_safe());
    }

    #[test]
    fn test_safe_describe() {
        assert!(validate_sql_safety("DESCRIBE users", 3).is_safe());
    }

    #[test]
    fn test_safe_explain() {
        assert!(validate_sql_safety("EXPLAIN SELECT * FROM users", 3).is_safe());
    }

    #[test]
    fn test_reject_empty() {
        assert!(!validate_sql_safety("", 3).is_safe());
    }

    #[test]
    fn test_reject_drop() {
        assert!(!validate_sql_safety("DROP TABLE users", 3).is_safe());
    }

    #[test]
    fn test_reject_delete() {
        assert!(!validate_sql_safety("DELETE FROM users WHERE id = 1", 3).is_safe());
    }

    #[test]
    fn test_reject_update() {
        assert!(!validate_sql_safety("UPDATE users SET name = 'x'", 3).is_safe());
    }

    #[test]
    fn test_reject_insert() {
        assert!(!validate_sql_safety("INSERT INTO users VALUES (1, 'x')", 3).is_safe());
    }

    #[test]
    fn test_reject_alter() {
        assert!(!validate_sql_safety("ALTER TABLE users ADD COLUMN age INT", 3).is_safe());
    }

    #[test]
    fn test_reject_truncate() {
        assert!(!validate_sql_safety("TRUNCATE users", 3).is_safe());
    }

    #[test]
    fn test_reject_select_with_drop() {
        // SELECT is allowed, but DROP inside is blocked
        assert!(!validate_sql_safety("SELECT * FROM users; DROP TABLE users", 3).is_safe());
    }

    #[test]
    fn test_reject_deep_subquery() {
        let sql = "SELECT * FROM (SELECT * FROM (SELECT * FROM (SELECT * FROM (SELECT 1))))";
        assert!(!validate_sql_safety(sql, 3).is_safe());
    }

    #[test]
    fn test_allow_moderate_subquery() {
        let sql = "SELECT * FROM (SELECT * FROM (SELECT * FROM users))";
        assert!(validate_sql_safety(sql, 3).is_safe());
    }

    #[test]
    fn test_keyword_in_string_literal_ok() {
        // DROP inside a string literal should not trigger rejection
        let sql = "SELECT * FROM users WHERE name = 'DROP TABLE'";
        assert!(validate_sql_safety(sql, 3).is_safe());
    }

    #[test]
    fn test_keyword_as_substring_ok() {
        // "UPDATED_AT" contains UPDATE but as part of a column name
        let sql = "SELECT UPDATED_AT FROM users";
        assert!(validate_sql_safety(sql, 3).is_safe());
    }

    // ── SQL extraction tests ─────────────────────────────────────────────

    #[test]
    fn test_extract_sql_from_code_fence() {
        let response = "Here is the SQL:\n```sql\nSELECT * FROM users\n```\nDone.";
        assert_eq!(extract_sql_from_response(response), "SELECT * FROM users");
    }

    #[test]
    fn test_extract_sql_from_generic_fence() {
        let response = "```\nSELECT id FROM orders\n```";
        assert_eq!(extract_sql_from_response(response), "SELECT id FROM orders");
    }

    #[test]
    fn test_extract_sql_no_fence() {
        let response = "SELECT count(*) FROM products WHERE price > 100";
        assert_eq!(
            extract_sql_from_response(response),
            "SELECT count(*) FROM products WHERE price > 100"
        );
    }

    #[test]
    fn test_extract_sql_with_explanation() {
        let response = "I'll query the users table:\nSELECT * FROM users LIMIT 10";
        assert_eq!(
            extract_sql_from_response(response),
            "SELECT * FROM users LIMIT 10"
        );
    }

    // ── LIMIT enforcement tests ──────────────────────────────────────────

    #[test]
    fn test_ensure_limit_adds_when_missing() {
        assert_eq!(
            ensure_limit("SELECT * FROM users", 1000),
            "SELECT * FROM users LIMIT 1000"
        );
    }

    #[test]
    fn test_ensure_limit_keeps_existing_small() {
        let sql = "SELECT * FROM users LIMIT 10";
        assert_eq!(ensure_limit(sql, 1000), sql);
    }

    #[test]
    fn test_ensure_limit_caps_large() {
        let sql = "SELECT * FROM users LIMIT 99999";
        assert_eq!(ensure_limit(sql, 1000), "SELECT * FROM users LIMIT 1000");
    }

    #[test]
    fn test_ensure_limit_skips_show() {
        assert_eq!(ensure_limit("SHOW TABLES", 1000), "SHOW TABLES");
    }

    #[test]
    fn test_ensure_limit_strips_semicolon() {
        assert_eq!(
            ensure_limit("SELECT * FROM users;", 1000),
            "SELECT * FROM users LIMIT 1000"
        );
    }

    // ── Schema context tests ─────────────────────────────────────────────

    #[test]
    fn test_schema_context_to_prompt() {
        let mut schema = SchemaContext::new();
        schema.add_table(TableInfo {
            name: "users".to_string(),
            columns: vec![
                ColumnInfo {
                    name: "id".to_string(),
                    data_type: "INT".to_string(),
                    nullable: false,
                },
                ColumnInfo {
                    name: "name".to_string(),
                    data_type: "VARCHAR".to_string(),
                    nullable: true,
                },
            ],
            row_count: Some(1000),
        });

        let prompt = schema.to_prompt_context();
        assert!(prompt.contains("users"));
        assert!(prompt.contains("id"));
        assert!(prompt.contains("name"));
        assert!(prompt.contains("~1000 rows"));
    }

    #[test]
    fn test_schema_context_empty() {
        let schema = SchemaContext::new();
        assert!(schema.is_empty());
        assert!(schema.to_prompt_context().contains("No tables"));
    }

    // ── Engine creation tests ────────────────────────────────────────────

    #[test]
    fn test_engine_creation() {
        let engine = NlToSqlEngine::new();
        assert_eq!(engine.config().max_retries, 3);
        assert!(engine.schema().is_empty());
        assert!(engine.history().is_empty());
    }

    #[test]
    fn test_engine_with_schema() {
        let mut engine = NlToSqlEngine::new();
        engine.add_table(TableInfo {
            name: "products".to_string(),
            columns: vec![ColumnInfo {
                name: "price".to_string(),
                data_type: "FLOAT".to_string(),
                nullable: false,
            }],
            row_count: None,
        });
        assert_eq!(engine.schema().tables.len(), 1);
    }

    // ── NlQueryResult tests ──────────────────────────────────────────────

    #[test]
    fn test_nl_query_result_serialization() {
        let result = NlQueryResult {
            question: "How many users?".to_string(),
            sql: "SELECT count(*) FROM users".to_string(),
            executed: true,
            data: vec![serde_json::json!({"count": 42})],
            row_count: 1,
            answer: Some("There are 42 users.".to_string()),
            explanation: None,
            execution_time_ms: 12.5,
            correction_rounds: 0,
            sql_history: Vec::new(),
        };

        let json = serde_json::to_string(&result).unwrap();
        assert!(json.contains("42"));
        assert!(json.contains("How many users?"));
    }

    // ── Strip fences tests ───────────────────────────────────────────────

    #[test]
    fn test_strip_sql_fences() {
        let sql = "```sql\nSELECT 1\n```";
        assert_eq!(strip_sql_fences(sql), "SELECT 1");
    }

    #[test]
    fn test_strip_no_fences() {
        let sql = "SELECT 1";
        assert_eq!(strip_sql_fences(sql), "SELECT 1");
    }

    // ── Subquery depth tests ─────────────────────────────────────────────

    #[test]
    fn test_subquery_depth_zero() {
        assert_eq!(count_subquery_depth("SELECT * FROM users"), 0);
    }

    #[test]
    fn test_subquery_depth_one() {
        assert_eq!(
            count_subquery_depth("SELECT * FROM (SELECT id FROM users)"),
            1
        );
    }

    #[test]
    fn test_subquery_depth_two() {
        assert_eq!(
            count_subquery_depth("SELECT * FROM (SELECT * FROM (SELECT 1))"),
            2
        );
    }

    #[test]
    fn test_subquery_depth_four() {
        assert_eq!(
            count_subquery_depth(
                "SELECT * FROM (SELECT * FROM (SELECT * FROM (SELECT * FROM (SELECT 1))))"
            ),
            4
        );
    }
}
