//! Integration tests for the SQL query engine.
//!
//! These tests verify:
//! - Lance dataset creation and SQL table registration
//! - Standard SQL: SELECT, WHERE, ORDER BY, LIMIT, GROUP BY, aggregations
//! - Custom UDFs: vector_distance, vector_norm, cosine_similarity, dot_product
//! - LanceTableProvider: cached and lazy scan modes
//! - SqlQueryEngine: high-level API

use std::sync::Arc;

use arrow::array::{
    Array, Float32Array, Float64Array, Int64Array, StringArray,
    FixedSizeListArray, RecordBatch,
};
use arrow_schema::{DataType, Field, Schema};
use tempfile::tempdir;

use crate::storage::lance::lance_backend::{LanceConfig, LanceStorageBackend};
use crate::storage::StorageBackend;
use crate::types::memory::MemoryNode;

use datafusion::catalog::TableProvider;

use super::context::AifastDbSessionContext;
use super::engine::{SqlEngineConfig, SqlQueryEngine};
use super::table_provider::LanceTableProvider;

// ============================================================================
// Helper: create a Lance dataset with test data
// ============================================================================

/// Create a Lance dataset with sample MemoryNodes and return its path.
fn create_test_lance_dataset(num_rows: usize, embedding_dim: usize) -> (tempfile::TempDir, String) {
    let dir = tempdir().unwrap();
    let lance_path = dir.path().join("test_sql.lance");

    let config = LanceConfig::new(&lance_path)
        .with_embedding_dim(embedding_dim as i32)
        .with_write_batch_size(num_rows); // flush in one batch

    let backend = LanceStorageBackend::open(config).unwrap();

    for i in 0..num_rows {
        let embedding: Vec<f32> = (0..embedding_dim)
            .map(|d| (i as f32 * 0.1) + (d as f32 * 0.01))
            .collect();

        let importance = (i as f64) / (num_rows as f64);

        let node = MemoryNode::new(format!("Document {}", i))
            .with_embedding(embedding)
            .with_importance(importance)
            .with_tag("test");

        backend.put(node).unwrap();
    }
    backend.flush().unwrap();

    let path_str = lance_path.to_string_lossy().to_string();
    (dir, path_str)
}

// ============================================================================
// Test: AifastDbSessionContext creation
// ============================================================================

#[test]
fn test_session_context_creation() {
    let ctx = AifastDbSessionContext::new();
    assert!(ctx.is_ok(), "Should create session context");
}

// ============================================================================
// Test: LanceTableProvider — cached mode
// ============================================================================

#[test]
fn test_lance_table_provider_cached() {
    let (_dir, path) = create_test_lance_dataset(10, 4);

    let rt = tokio::runtime::Runtime::new().unwrap();
    let provider = rt.block_on(async {
        LanceTableProvider::try_new(&path, None).await
    });

    assert!(provider.is_ok(), "Should open Lance dataset: {:?}", provider.err());
    let provider = provider.unwrap();

    assert_eq!(provider.num_rows(), 10);

    let schema = provider.schema();
    assert!(schema.field_with_name("content").is_ok(), "Should have 'content' field");
    assert!(schema.field_with_name("embedding").is_ok(), "Should have 'embedding' field");
}

// ============================================================================
// Test: LanceTableProvider — lazy scan mode
// ============================================================================

#[test]
fn test_lance_table_provider_lazy() {
    let (_dir, path) = create_test_lance_dataset(10, 4);

    let rt = tokio::runtime::Runtime::new().unwrap();
    let provider = rt.block_on(async {
        LanceTableProvider::try_new_lazy(&path, None).await
    });

    assert!(provider.is_ok(), "Should open Lance dataset lazy: {:?}", provider.err());
    let provider = provider.unwrap();

    // Lazy mode doesn't cache, so num_rows() should be 0
    assert_eq!(provider.num_rows(), 0);

    // But schema should be valid
    let schema = provider.schema();
    assert!(schema.fields().len() > 0, "Should have fields in schema");
}

// ============================================================================
// Test: SqlQueryEngine — register and basic SELECT
// ============================================================================

#[test]
fn test_sql_engine_basic_select() {
    let (_dir, path) = create_test_lance_dataset(5, 4);

    let engine = SqlQueryEngine::new().unwrap();
    engine.register_table("docs", &path).unwrap();

    // Verify table is registered
    let tables = engine.tables();
    assert!(tables.contains(&"docs".to_string()), "Table 'docs' should be registered");

    // SELECT *
    let result = engine.query("SELECT * FROM docs");
    assert!(result.is_ok(), "SELECT * should work: {:?}", result.err());

    let batches = result.unwrap();
    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
    assert_eq!(total_rows, 5, "Should have 5 rows");
}

// ============================================================================
// Test: SELECT with WHERE filter
// ============================================================================

#[test]
fn test_sql_where_filter() {
    let (_dir, path) = create_test_lance_dataset(10, 4);

    let engine = SqlQueryEngine::new().unwrap();
    engine.register_table("docs", &path).unwrap();

    // Filter with WHERE importance > 0.5
    let result = engine.query("SELECT content, importance FROM docs WHERE importance > 0.5");
    assert!(result.is_ok(), "WHERE filter should work: {:?}", result.err());

    let batches = result.unwrap();
    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
    assert!(total_rows > 0, "Should have some rows with importance > 0.5");
    assert!(total_rows < 10, "Should have fewer rows after filtering");
}

// ============================================================================
// Test: ORDER BY + LIMIT
// ============================================================================

#[test]
fn test_sql_order_by_limit() {
    let (_dir, path) = create_test_lance_dataset(10, 4);

    let engine = SqlQueryEngine::new().unwrap();
    engine.register_table("docs", &path).unwrap();

    let result = engine.query(
        "SELECT content, importance FROM docs ORDER BY importance DESC LIMIT 3"
    );
    assert!(result.is_ok(), "ORDER BY LIMIT should work: {:?}", result.err());

    let batches = result.unwrap();
    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
    assert_eq!(total_rows, 3, "LIMIT 3 should return 3 rows");

    // Verify ordering: first row should have highest importance
    if let Some(batch) = batches.first() {
        let importance_col = batch.column_by_name("importance");
        assert!(importance_col.is_some(), "Should have importance column");

        if let Some(col) = importance_col {
            if let Some(arr) = col.as_any().downcast_ref::<Float64Array>() {
                if arr.len() >= 2 {
                    assert!(
                        arr.value(0) >= arr.value(1),
                        "Results should be ordered DESC: {} >= {}",
                        arr.value(0),
                        arr.value(1)
                    );
                }
            }
        }
    }
}

// ============================================================================
// Test: COUNT aggregation
// ============================================================================

#[test]
fn test_sql_count_aggregation() {
    let (_dir, path) = create_test_lance_dataset(10, 4);

    let engine = SqlQueryEngine::new().unwrap();
    engine.register_table("docs", &path).unwrap();

    let result = engine.query("SELECT COUNT(*) AS cnt FROM docs");
    assert!(result.is_ok(), "COUNT should work: {:?}", result.err());

    let batches = result.unwrap();
    assert!(!batches.is_empty(), "Should have result batch");

    let batch = &batches[0];
    let cnt_col = batch.column(0);
    let cnt = cnt_col
        .as_any()
        .downcast_ref::<Int64Array>()
        .expect("COUNT should return Int64");
    assert_eq!(cnt.value(0), 10, "COUNT should be 10");
}

// ============================================================================
// Test: AVG, MIN, MAX aggregation
// ============================================================================

#[test]
fn test_sql_aggregations() {
    let (_dir, path) = create_test_lance_dataset(10, 4);

    let engine = SqlQueryEngine::new().unwrap();
    engine.register_table("docs", &path).unwrap();

    let result = engine.query(
        "SELECT AVG(importance) AS avg_imp, MIN(importance) AS min_imp, MAX(importance) AS max_imp FROM docs"
    );
    assert!(result.is_ok(), "Aggregations should work: {:?}", result.err());

    let batches = result.unwrap();
    assert!(!batches.is_empty());

    let batch = &batches[0];
    assert_eq!(batch.num_rows(), 1, "Aggregation should return 1 row");
    assert_eq!(batch.num_columns(), 3, "Should have 3 aggregation columns");
}

// ============================================================================
// Test: EXPLAIN
// ============================================================================

#[test]
fn test_sql_explain() {
    let (_dir, path) = create_test_lance_dataset(5, 4);

    let engine = SqlQueryEngine::new().unwrap();
    engine.register_table("docs", &path).unwrap();

    let result = engine.explain("SELECT * FROM docs WHERE importance > 0.5");
    assert!(result.is_ok(), "EXPLAIN should work: {:?}", result.err());

    let plan = result.unwrap();
    assert!(!plan.is_empty(), "EXPLAIN should return a plan");
}

// ============================================================================
// Test: query_to_string (pretty print)
// ============================================================================

#[test]
fn test_query_to_string() {
    let (_dir, path) = create_test_lance_dataset(3, 4);

    let engine = SqlQueryEngine::new().unwrap();
    engine.register_table("docs", &path).unwrap();

    let result = engine.query_to_string("SELECT content, importance FROM docs LIMIT 2");
    assert!(result.is_ok(), "query_to_string should work: {:?}", result.err());

    let output = result.unwrap();
    assert!(output.contains("content"), "Output should contain 'content' header");
    assert!(output.contains("importance"), "Output should contain 'importance' header");
}

// ============================================================================
// Test: query_count
// ============================================================================

#[test]
fn test_query_count() {
    let (_dir, path) = create_test_lance_dataset(7, 4);

    let engine = SqlQueryEngine::new().unwrap();
    engine.register_table("docs", &path).unwrap();

    let count = engine.query_count("SELECT * FROM docs WHERE importance > 0.3");
    assert!(count.is_ok(), "query_count should work: {:?}", count.err());

    let cnt = count.unwrap();
    assert!(cnt > 0, "Should have some rows");
    assert!(cnt <= 7, "Should not exceed total rows");
}

// ============================================================================
// Test: SqlEngineConfig with_table
// ============================================================================

#[test]
fn test_sql_engine_config() {
    let (_dir, path) = create_test_lance_dataset(5, 4);

    let config = SqlEngineConfig::new()
        .with_table("my_docs", &path);

    let engine = SqlQueryEngine::with_config(config);
    assert!(engine.is_ok(), "Config-based creation should work: {:?}", engine.err());

    let engine = engine.unwrap();
    let result = engine.query("SELECT COUNT(*) FROM my_docs");
    assert!(result.is_ok(), "Query on config-registered table should work");
}

// ============================================================================
// Test: Multiple tables
// ============================================================================

#[test]
fn test_multiple_tables() {
    let (_dir1, path1) = create_test_lance_dataset(5, 4);
    let (_dir2, path2) = create_test_lance_dataset(8, 4);

    let engine = SqlQueryEngine::new().unwrap();
    engine.register_table("table_a", &path1).unwrap();
    engine.register_table("table_b", &path2).unwrap();

    let tables = engine.tables();
    assert!(tables.contains(&"table_a".to_string()));
    assert!(tables.contains(&"table_b".to_string()));

    // Query each table
    let count_a = engine.query_count("SELECT * FROM table_a").unwrap();
    let count_b = engine.query_count("SELECT * FROM table_b").unwrap();

    assert_eq!(count_a, 5);
    assert_eq!(count_b, 8);
}

// ============================================================================
// Test: SELECT with column projection
// ============================================================================

#[test]
fn test_column_projection() {
    let (_dir, path) = create_test_lance_dataset(5, 4);

    let engine = SqlQueryEngine::new().unwrap();
    engine.register_table("docs", &path).unwrap();

    let result = engine.query("SELECT content FROM docs LIMIT 2");
    assert!(result.is_ok());

    let batches = result.unwrap();
    assert!(!batches.is_empty());

    let batch = &batches[0];
    // Should only have 'content' column
    assert_eq!(batch.num_columns(), 1, "Projection should return 1 column");
    assert_eq!(
        batch.schema().field(0).name(),
        "content",
        "Projected column should be 'content'"
    );
}

// ============================================================================
// Test: sql_schema
// ============================================================================

#[test]
fn test_sql_schema() {
    let (_dir, path) = create_test_lance_dataset(3, 4);

    let engine = SqlQueryEngine::new().unwrap();
    engine.register_table("docs", &path).unwrap();

    let schema = engine.query_schema("SELECT content, importance FROM docs");
    assert!(schema.is_ok(), "query_schema should work: {:?}", schema.err());

    let schema = schema.unwrap();
    assert_eq!(schema.fields().len(), 2);
    assert_eq!(schema.field(0).name(), "content");
    assert_eq!(schema.field(1).name(), "importance");
}

// ============================================================================
// Test: UDF — vector_norm (with real data from Lance)
// ============================================================================

#[test]
fn test_udf_registration() {
    // Just verify that UDFs are registered and the engine can parse queries using them
    let ctx = AifastDbSessionContext::new().unwrap();

    // These UDFs should be registered
    let _session = ctx.session();

    // Check UDF names are registered (they should be callable)
    // We test by running a simple SELECT with scalar arguments
    let result = ctx.execute_sql("SELECT 1 + 1 AS sum");
    assert!(result.is_ok(), "Simple SQL should work: {:?}", result.err());
}

// ============================================================================
// Test: Empty result set
// ============================================================================

#[test]
fn test_empty_result() {
    let (_dir, path) = create_test_lance_dataset(5, 4);

    let engine = SqlQueryEngine::new().unwrap();
    engine.register_table("docs", &path).unwrap();

    // Filter that matches nothing
    let result = engine.query("SELECT * FROM docs WHERE importance > 999.0");
    assert!(result.is_ok());

    let batches = result.unwrap();
    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
    assert_eq!(total_rows, 0, "Should have no matching rows");
}

// ============================================================================
// Test: DISTINCT
// ============================================================================

#[test]
fn test_sql_distinct() {
    let (_dir, path) = create_test_lance_dataset(10, 4);

    let engine = SqlQueryEngine::new().unwrap();
    engine.register_table("docs", &path).unwrap();

    // All documents have tag "test"
    let result = engine.query("SELECT DISTINCT content_type FROM docs");
    assert!(result.is_ok(), "DISTINCT should work: {:?}", result.err());
}
