use datafusion::arrow::array::{ArrayRef, Int64Array, StringArray};
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::catalog::{TableFunctionImpl, TableProvider};
use datafusion::common::{Result, ScalarValue, plan_err};
use datafusion::datasource::memory::MemTable;
use datafusion::logical_expr::Expr;
use datafusion::prelude::SessionContext;
use std::sync::Arc;
 
#[derive(Debug)]
pub struct TransformFunction {}
 
impl TableFunctionImpl for TransformFunction {
    fn call(&self, exprs: &[Expr]) -> Result<Arc<dyn TableProvider>> {
        if exprs.len() != 3 {
            return plan_err!(
                "Expected exactly three arguments: a, b, c, but got {}",
                exprs.len()
            );
        }
 
        println!("{:?}", exprs);
 
        let extract_int64 = |expr: &Expr, arg_name: &str| -> Result<i64> {
            match expr {
                Expr::Literal(ScalarValue::Int64(Some(val)), _) => Ok(*val),
                // Expr::Column()
                _ => plan_err!("Argument {} must be an Int64 literal", arg_name),
            }
        };
 
        let a = extract_int64(&exprs[0], "a")?;
        let b = extract_int64(&exprs[1], "b")?;
        let c = extract_int64(&exprs[2], "c")?;
 
        // Compute output columns: x = a + b, y = b * c
        let x = a + b;
        let y = b * c;
 
        // Define output schema
        let schema = Arc::new(Schema::new(vec![
            Field::new("x", DataType::Int64, false),
            Field::new("y", DataType::Int64, false),
        ]));
 
        // Create output arrays
        let x_array = Arc::new(Int64Array::from(vec![x])) as ArrayRef;
        let y_array = Arc::new(Int64Array::from(vec![y])) as ArrayRef;
 
        // Create a single RecordBatch
        let batch = RecordBatch::try_new(schema.clone(), vec![x_array, y_array])?;
 
        // Wrap in a MemTable
        let provider = MemTable::try_new(schema, vec![vec![batch]])?;
 
        Ok(Arc::new(provider))
    }
}
 
// --- Usage Example ---
 
// /// Registers the TransformFunction as a TableUDF in the SessionContext.
fn register_udtf(ctx: &mut SessionContext) -> Result<()> {
    // 1. Create the implementation instance
    let udtf = Arc::new(TransformFunction {});
    ctx.register_udtf("my_transform", udtf);
 
    Ok(())
}
 
/// Creates a small in-memory table for demonstration.
fn create_dummy_table(ctx: &mut SessionContext) -> Result<()> {
    let schema = Arc::new(Schema::new(vec![
        Field::new("id", DataType::Utf8, false),
        Field::new("a", DataType::Int64, false),
        Field::new("b", DataType::Int64, false),
        Field::new("c", DataType::Int64, false),
    ]));
 
    let batch = RecordBatch::try_new(
        schema.clone(),
        vec![
            Arc::new(StringArray::from(vec!["r1", "r2"])) as ArrayRef,
            Arc::new(Int64Array::from(vec![10, 20])) as ArrayRef,
            Arc::new(Int64Array::from(vec![5, 6])) as ArrayRef,
            Arc::new(Int64Array::from(vec![2, 3])) as ArrayRef,
        ],
    )?;
 
    let provider = MemTable::try_new(schema, vec![vec![batch]])?;
    ctx.register_table("my_table", Arc::new(provider))?;
    Ok(())
}
 
#[tokio::main]
async fn main() -> Result<()> {
    let mut ctx = SessionContext::new();
 
    // 1. Register the custom UDTF
    register_udtf(&mut ctx)?;
 
    // 2. Register a dummy table
    create_dummy_table(&mut ctx)?;
 
    // 3. Define and execute the SQL query
    let sql = r#"
        SELECT 
            t1.id, 
            t2.x AS a_plus_b, 
            t2.y AS b_times_c
        FROM 
            my_table AS t1,
            LATERAL my_transform(1, 2, 3) AS t2(x, y)
    "#;
 
    // let sql = r#"
    //     SELECT 
    //         t1.id, 
    //         t2.x AS a_plus_b, 
    //         t2.y AS b_times_c
    //     FROM 
    //         my_table AS t1,
    //         LATERAL my_transform(t1.a, t1.b, t1.c) AS t2(x, y)
    // "#;
 
 
    println!("Executing SQL:\n{}", sql);
 
    let df = ctx.sql(sql).await?;
 
    println!("\nQuery Result:");
    df.show().await?;
 
    Ok(())
}