//! SqlQueryEngine — unified SQL query entry point for AiFastDb.
//!
//! Wraps `AifastDbSessionContext` and provides a high-level API for
//! executing SQL queries against registered Lance tables.

use std::collections::HashMap;

use arrow::array::RecordBatch;
use arrow_schema::SchemaRef;

use super::context::{AifastDbSessionContext, SqlResult};

/// Configuration for the SQL query engine.
#[derive(Debug, Clone, Default)]
pub struct SqlEngineConfig {
    /// Tables to auto-register on creation: (name, path/uri)
    pub tables: Vec<(String, String)>,
    /// Storage options for remote backends
    pub storage_options: Option<HashMap<String, String>>,
}

impl SqlEngineConfig {
    pub fn new() -> Self {
        Self::default()
    }

    /// Add a table to auto-register.
    pub fn with_table(mut self, name: impl Into<String>, path: impl Into<String>) -> Self {
        self.tables.push((name.into(), path.into()));
        self
    }

    /// Set storage options for remote backends.
    pub fn with_storage_options(mut self, opts: HashMap<String, String>) -> Self {
        self.storage_options = Some(opts);
        self
    }
}

/// High-level SQL query engine for AiFastDb.
///
/// Provides a simple API for:
/// - Registering Lance datasets as SQL tables
/// - Executing SQL queries and getting results as Arrow RecordBatches
/// - Inspecting table schemas
///
/// # Example
///
/// ```ignore
/// use aifastdb_core::sql::SqlQueryEngine;
///
/// let engine = SqlQueryEngine::new()?;
/// engine.register_table("memories", "/data/memories.lance")?;
///
/// let results = engine.query(
///     "SELECT id, content, importance FROM memories
///      WHERE importance > 0.8
///      ORDER BY importance DESC
///      LIMIT 10"
/// )?;
/// ```
pub struct SqlQueryEngine {
    ctx: AifastDbSessionContext,
}

impl SqlQueryEngine {
    /// Create a new SQL query engine with default settings.
    pub fn new() -> SqlResult<Self> {
        let ctx = AifastDbSessionContext::new()?;
        Ok(Self { ctx })
    }

    /// Create with configuration (auto-registers tables).
    pub fn with_config(config: SqlEngineConfig) -> SqlResult<Self> {
        let ctx = AifastDbSessionContext::new()?;

        for (name, path) in &config.tables {
            if let Some(ref opts) = config.storage_options {
                ctx.register_lance_table_with_options(name, path, opts.clone())?;
            } else {
                ctx.register_lance_table(name, path)?;
            }
        }

        Ok(Self { ctx })
    }

    /// Register a Lance dataset as a named table.
    pub fn register_table(&self, name: &str, path: &str) -> SqlResult<()> {
        self.ctx.register_lance_table(name, path)
    }

    /// Register a Lance dataset with custom storage options.
    pub fn register_table_with_options(
        &self,
        name: &str,
        uri: &str,
        storage_options: HashMap<String, String>,
    ) -> SqlResult<()> {
        self.ctx
            .register_lance_table_with_options(name, uri, storage_options)
    }

    /// Execute a SQL query and return results as Arrow RecordBatches.
    pub fn query(&self, sql: &str) -> SqlResult<Vec<RecordBatch>> {
        self.ctx.execute_sql(sql)
    }

    /// Execute a SQL query and return the result schema.
    pub fn query_schema(&self, sql: &str) -> SqlResult<SchemaRef> {
        self.ctx.sql_schema(sql)
    }

    /// List all registered table names.
    pub fn tables(&self) -> Vec<String> {
        self.ctx.table_names()
    }

    /// Execute a SQL query and return results as a formatted string table.
    /// Useful for debugging and CLI output.
    pub fn query_to_string(&self, sql: &str) -> SqlResult<String> {
        let batches = self.query(sql)?;

        if batches.is_empty() {
            return Ok("(empty result set)".to_string());
        }

        // Use Arrow's built-in pretty printer
        let formatted = arrow::util::pretty::pretty_format_batches(&batches)
            .map_err(|e| super::context::SqlError::Lance(format!("Format error: {}", e)))?;

        Ok(formatted.to_string())
    }

    /// Get the query execution plan as a string (EXPLAIN output).
    ///
    /// Useful for understanding how DataFusion will execute the query
    /// and which optimizations (projection/filter push-down) are applied.
    pub fn explain(&self, sql: &str) -> SqlResult<String> {
        let explain_sql = format!("EXPLAIN {}", sql);
        self.query_to_string(&explain_sql)
    }

    /// Get the verbose query execution plan (EXPLAIN VERBOSE).
    pub fn explain_verbose(&self, sql: &str) -> SqlResult<String> {
        let explain_sql = format!("EXPLAIN VERBOSE {}", sql);
        self.query_to_string(&explain_sql)
    }

    /// Execute a SQL query and return the total row count.
    pub fn query_count(&self, sql: &str) -> SqlResult<usize> {
        let batches = self.query(sql)?;
        Ok(batches.iter().map(|b| b.num_rows()).sum())
    }

    /// Get the underlying AifastDbSessionContext.
    pub fn session_context(&self) -> &AifastDbSessionContext {
        &self.ctx
    }

    /// Get a mutable reference to the session context.
    pub fn session_context_mut(&mut self) -> &mut AifastDbSessionContext {
        &mut self.ctx
    }
}
