//! AiFastDb SessionContext — wraps DataFusion's SessionContext with
//! pre-registered Lance tables, custom UDFs, and helper methods.

use std::collections::HashMap;
use std::sync::Arc;

use arrow_schema::SchemaRef;
use datafusion::execution::context::SessionContext;
use datafusion::error::DataFusionError;
use datafusion::arrow::record_batch::RecordBatch;

use super::table_provider::LanceTableProvider;
use super::udf::register_all_udfs;

/// Error type for SQL operations.
#[derive(Debug, thiserror::Error)]
pub enum SqlError {
    #[error("DataFusion error: {0}")]
    DataFusion(#[from] DataFusionError),

    #[error("Lance error: {0}")]
    Lance(String),

    #[error("Table not found: {0}")]
    TableNotFound(String),

    #[error("Configuration error: {0}")]
    Config(String),
}

pub type SqlResult<T> = Result<T, SqlError>;

/// A wrapper around DataFusion `SessionContext` that provides AiFastDb-specific
/// functionality, including automatic Lance table registration, custom UDFs,
/// and convenience methods for SQL execution.
///
/// # Example
///
/// ```ignore
/// use aifastdb_core::sql::AifastDbSessionContext;
///
/// let ctx = AifastDbSessionContext::new();
/// ctx.register_lance_table("memories", "/path/to/dataset.lance").await?;
/// let results = ctx.sql("SELECT id, content FROM memories WHERE importance > 0.5").await?;
/// ```
pub struct AifastDbSessionContext {
    /// The underlying DataFusion SessionContext.
    inner: SessionContext,
    /// Tokio runtime for async operations in sync context.
    runtime: tokio::runtime::Runtime,
}

impl AifastDbSessionContext {
    /// Create a new AiFastDb session context with default configuration.
    ///
    /// This automatically registers all custom UDFs (vector_distance, gpt_*, etc.).
    pub fn new() -> SqlResult<Self> {
        let ctx = SessionContext::new();

        // Register all custom UDFs (vector operations)
        register_all_udfs(&ctx)?;

        // Register GPT UDFs when candle feature is enabled
        #[cfg(feature = "candle")]
        super::gpt_udf::register_gpt_udfs(&ctx);

        let runtime = tokio::runtime::Builder::new_multi_thread()
            .enable_all()
            .build()
            .map_err(|e| SqlError::Config(format!("Failed to create tokio runtime: {}", e)))?;

        Ok(Self {
            inner: ctx,
            runtime,
        })
    }

    /// Create from an existing DataFusion SessionContext.
    pub fn from_session(ctx: SessionContext) -> SqlResult<Self> {
        register_all_udfs(&ctx)?;

        // Register GPT UDFs when candle feature is enabled
        #[cfg(feature = "candle")]
        super::gpt_udf::register_gpt_udfs(&ctx);

        let runtime = tokio::runtime::Builder::new_multi_thread()
            .enable_all()
            .build()
            .map_err(|e| SqlError::Config(format!("Failed to create tokio runtime: {}", e)))?;

        Ok(Self {
            inner: ctx,
            runtime,
        })
    }

    /// Register a Lance dataset as a named table in this session.
    ///
    /// The table becomes available for SQL queries using the given name.
    /// Supports both local paths and remote URIs (s3://, gs://, az://).
    ///
    /// # Arguments
    /// * `table_name` - The name to use in SQL queries (e.g., "memories")
    /// * `dataset_path` - Path or URI to the Lance dataset
    pub fn register_lance_table(
        &self,
        table_name: &str,
        dataset_path: &str,
    ) -> SqlResult<()> {
        let provider = self.runtime.block_on(async {
            LanceTableProvider::try_new(dataset_path, None).await
        })?;

        self.inner
            .register_table(table_name, Arc::new(provider))
            .map_err(SqlError::DataFusion)?;

        Ok(())
    }

    /// Register a Lance dataset with custom storage options (for remote URIs).
    pub fn register_lance_table_with_options(
        &self,
        table_name: &str,
        dataset_uri: &str,
        storage_options: HashMap<String, String>,
    ) -> SqlResult<()> {
        let provider = self.runtime.block_on(async {
            LanceTableProvider::try_new(dataset_uri, Some(storage_options)).await
        })?;

        self.inner
            .register_table(table_name, Arc::new(provider))
            .map_err(SqlError::DataFusion)?;

        Ok(())
    }

    /// Execute a SQL query and return the results as Arrow RecordBatches.
    pub fn execute_sql(&self, sql: &str) -> SqlResult<Vec<RecordBatch>> {
        self.runtime.block_on(async {
            let df = self.inner.sql(sql).await?;
            let batches = df.collect().await?;
            Ok(batches)
        })
    }

    /// Execute a SQL query and return the schema of the result.
    pub fn sql_schema(&self, sql: &str) -> SqlResult<SchemaRef> {
        self.runtime.block_on(async {
            let df = self.inner.sql(sql).await?;
            let schema = df.schema();
            // Convert DFSchema to Arrow Schema
            let arrow_schema: arrow_schema::Schema = schema.as_arrow().as_ref().clone();
            Ok(Arc::new(arrow_schema))
        })
    }

    /// List all registered table names.
    pub fn table_names(&self) -> Vec<String> {
        self.inner
            .catalog_names()
            .iter()
            .flat_map(|catalog_name| {
                let catalog = self.inner.catalog(catalog_name);
                match catalog {
                    Some(catalog) => catalog
                        .schema_names()
                        .iter()
                        .flat_map(|schema_name| match catalog.schema(schema_name) {
                            Some(schema) => schema.table_names(),
                            None => Vec::new(),
                        })
                        .collect::<Vec<_>>(),
                    None => Vec::new(),
                }
            })
            .collect()
    }

    /// Get a reference to the underlying DataFusion SessionContext.
    pub fn session(&self) -> &SessionContext {
        &self.inner
    }

    /// Get a mutable reference to the underlying DataFusion SessionContext.
    pub fn session_mut(&mut self) -> &mut SessionContext {
        &mut self.inner
    }
}

impl Default for AifastDbSessionContext {
    fn default() -> Self {
        Self::new().expect("Failed to create default AifastDbSessionContext")
    }
}
