use ducklake_catalog::{SnapshotId, SortExpression, SortId, SortInfo, TableId};
use sqlparser::ast::{Expr, Ident, VisitMut, VisitorMut};
use sqlparser::dialect::DuckDbDialect;
use sqlparser::parser::Parser;
use sqlparser::tokenizer::Token;
use std::ops::ControlFlow;

use crate::{Result, StorageError};

#[cfg(any(
    feature = "sqlite",
    feature = "duckdb",
    feature = "postgres",
    feature = "quack"
))]
pub(crate) fn rewrite_sort_expression_for_renamed_column(
    expression: &str,
    old_column_name: &str,
    new_column_name: &str,
) -> Result<String> {
    let dialect = DuckDbDialect {};
    let mut parser = Parser::new(&dialect)
        .try_with_sql(expression)
        .map_err(|err| {
            StorageError::Backend(format!("DuckLake sort expression parse error: {err}"))
        })?;
    let mut expression = parser.parse_expr().map_err(|err| {
        StorageError::Backend(format!("DuckLake sort expression parse error: {err}"))
    })?;
    parser.expect_token(&Token::EOF).map_err(|err| {
        StorageError::Backend(format!("DuckLake sort expression parse error: {err}"))
    })?;

    struct RenameColumnRef<'a> {
        old_column_name: &'a str,
        new_column_name: &'a str,
    }

    impl VisitorMut for RenameColumnRef<'_> {
        type Break = ();

        fn post_visit_expr(&mut self, expression: &mut Expr) -> ControlFlow<Self::Break> {
            if let Expr::Identifier(ident) = expression
                && ident.value.eq_ignore_ascii_case(self.old_column_name)
            {
                *ident = sort_expression_ident(self.new_column_name);
            }
            ControlFlow::Continue(())
        }
    }

    let mut visitor = RenameColumnRef {
        old_column_name,
        new_column_name,
    };
    let _ = expression.visit(&mut visitor);
    Ok(expression.to_string())
}

#[cfg(any(
    feature = "sqlite",
    feature = "duckdb",
    feature = "postgres",
    feature = "quack"
))]
pub(crate) fn sort_expression_references_unqualified_column(
    expression: &str,
    column_name: &str,
) -> Result<bool> {
    Ok(sort_expression_unqualified_column_refs(expression)?
        .iter()
        .any(|reference| reference.eq_ignore_ascii_case(column_name)))
}

#[cfg(any(
    feature = "sqlite",
    feature = "duckdb",
    feature = "postgres",
    feature = "quack"
))]
pub(crate) fn sort_expression_unqualified_column_refs(expression: &str) -> Result<Vec<String>> {
    let dialect = DuckDbDialect {};
    let mut parser = Parser::new(&dialect)
        .try_with_sql(expression)
        .map_err(|err| {
            StorageError::Backend(format!("DuckLake sort expression parse error: {err}"))
        })?;
    let mut expression = parser.parse_expr().map_err(|err| {
        StorageError::Backend(format!("DuckLake sort expression parse error: {err}"))
    })?;
    parser.expect_token(&Token::EOF).map_err(|err| {
        StorageError::Backend(format!("DuckLake sort expression parse error: {err}"))
    })?;

    struct FindColumnRefs {
        refs: Vec<String>,
    }

    impl VisitorMut for FindColumnRefs {
        type Break = ();

        fn post_visit_expr(&mut self, expression: &mut Expr) -> ControlFlow<Self::Break> {
            if let Expr::Identifier(ident) = expression {
                let already_seen = self
                    .refs
                    .iter()
                    .any(|reference| reference.eq_ignore_ascii_case(&ident.value));
                if !already_seen {
                    self.refs.push(ident.value.clone());
                }
            }
            ControlFlow::Continue(())
        }
    }

    let mut visitor = FindColumnRefs { refs: Vec::new() };
    let _ = expression.visit(&mut visitor);
    Ok(visitor.refs)
}

#[cfg(any(
    feature = "sqlite",
    feature = "duckdb",
    feature = "postgres",
    feature = "quack"
))]
fn sort_expression_ident(column_name: &str) -> Ident {
    if column_name
        .chars()
        .next()
        .is_some_and(|ch| ch == '_' || ch.is_ascii_alphabetic())
        && column_name
            .chars()
            .all(|ch| ch == '_' || ch.is_ascii_alphanumeric())
    {
        Ident::new(column_name)
    } else {
        Ident::with_quote('"', column_name)
    }
}

pub(crate) struct SortRenameRewrite<'a> {
    pub(crate) table_id: TableId,
    pub(crate) latest_snapshot: SnapshotId,
    pub(crate) new_snapshot: i64,
    pub(crate) next_catalog_id: i64,
    pub(crate) old_column_name: &'a str,
    pub(crate) new_column_name: &'a str,
}

#[cfg(any(feature = "sqlite", feature = "duckdb", test))]
pub(crate) fn rewrite_sort_expressions_for_rename_sync<
    TableExists,
    QueryInfos,
    QueryExpressions,
    CloseSort,
    ExecuteQuery,
>(
    rewrite: SortRenameRewrite<'_>,
    mut sort_table_exists: TableExists,
    mut query_sort_infos: QueryInfos,
    mut query_sort_expressions: QueryExpressions,
    mut close_sort: CloseSort,
    mut execute_sort_query: ExecuteQuery,
) -> Result<i64>
where
    TableExists: FnMut(&str) -> Result<bool>,
    QueryInfos: FnMut(crate::catalog_queries::CatalogQuery) -> Result<Vec<SortInfo>>,
    QueryExpressions: FnMut(crate::catalog_queries::CatalogQuery) -> Result<Vec<SortExpression>>,
    CloseSort: FnMut(SortId, TableId, i64) -> Result<usize>,
    ExecuteQuery: FnMut(crate::catalog_queries::CatalogQuery) -> Result<()>,
{
    if !sort_table_exists("ducklake_sort_info")? || !sort_table_exists("ducklake_sort_expression")?
    {
        return Ok(rewrite.next_catalog_id);
    }

    let active_sorts = query_sort_infos(crate::catalog_queries::active_sort_infos(
        rewrite.table_id,
        rewrite.latest_snapshot,
    ))?;

    let mut allocated = 0;
    for sort in active_sorts {
        let old_sort_id = sort.sort_id;
        let new_sort_id = SortId(rewrite.next_catalog_id + allocated);
        allocated += 1;

        validate_closed_sort_count(
            rewrite.table_id,
            old_sort_id,
            close_sort(old_sort_id, rewrite.table_id, rewrite.new_snapshot)?,
        )?;

        execute_sort_query(crate::catalog_queries::insert_sort_info(
            new_sort_id,
            rewrite.table_id,
            rewrite.new_snapshot,
        ))?;

        let expressions = query_sort_expressions(crate::catalog_queries::sort_expressions(
            old_sort_id,
            rewrite.table_id,
        ))?;
        for expression in expressions {
            execute_sort_query(crate::catalog_queries::insert_sort_expression_row(
                &rewritten_sort_expression(
                    expression,
                    new_sort_id,
                    rewrite.old_column_name,
                    rewrite.new_column_name,
                )?,
            ))?;
        }
    }

    Ok(rewrite.next_catalog_id + allocated)
}

#[cfg(any(feature = "postgres", feature = "quack"))]
pub(crate) async fn rewrite_sort_expressions_for_rename_async<
    TableExists,
    TableExistsFut,
    QueryInfos,
    QueryInfosFut,
    QueryExpressions,
    QueryExpressionsFut,
    CloseSort,
    CloseSortFut,
    ExecuteQuery,
    ExecuteQueryFut,
>(
    rewrite: SortRenameRewrite<'_>,
    mut sort_table_exists: TableExists,
    mut query_sort_infos: QueryInfos,
    mut query_sort_expressions: QueryExpressions,
    mut close_sort: CloseSort,
    mut execute_sort_query: ExecuteQuery,
) -> Result<i64>
where
    TableExists: FnMut(String) -> TableExistsFut,
    TableExistsFut: std::future::Future<Output = Result<bool>>,
    QueryInfos: FnMut(crate::catalog_queries::CatalogQuery) -> QueryInfosFut,
    QueryInfosFut: std::future::Future<Output = Result<Vec<SortInfo>>>,
    QueryExpressions: FnMut(crate::catalog_queries::CatalogQuery) -> QueryExpressionsFut,
    QueryExpressionsFut: std::future::Future<Output = Result<Vec<SortExpression>>>,
    CloseSort: FnMut(SortId, TableId, i64) -> CloseSortFut,
    CloseSortFut: std::future::Future<Output = Result<usize>>,
    ExecuteQuery: FnMut(crate::catalog_queries::CatalogQuery) -> ExecuteQueryFut,
    ExecuteQueryFut: std::future::Future<Output = Result<()>>,
{
    if !sort_table_exists("ducklake_sort_info".to_string()).await?
        || !sort_table_exists("ducklake_sort_expression".to_string()).await?
    {
        return Ok(rewrite.next_catalog_id);
    }

    let active_sorts = query_sort_infos(crate::catalog_queries::active_sort_infos(
        rewrite.table_id,
        rewrite.latest_snapshot,
    ))
    .await?;

    let mut allocated = 0;
    for sort in active_sorts {
        let old_sort_id = sort.sort_id;
        let new_sort_id = SortId(rewrite.next_catalog_id + allocated);
        allocated += 1;

        validate_closed_sort_count(
            rewrite.table_id,
            old_sort_id,
            close_sort(old_sort_id, rewrite.table_id, rewrite.new_snapshot).await?,
        )?;

        execute_sort_query(crate::catalog_queries::insert_sort_info(
            new_sort_id,
            rewrite.table_id,
            rewrite.new_snapshot,
        ))
        .await?;

        let expressions = query_sort_expressions(crate::catalog_queries::sort_expressions(
            old_sort_id,
            rewrite.table_id,
        ))
        .await?;
        for expression in expressions {
            execute_sort_query(crate::catalog_queries::insert_sort_expression_row(
                &rewritten_sort_expression(
                    expression,
                    new_sort_id,
                    rewrite.old_column_name,
                    rewrite.new_column_name,
                )?,
            ))
            .await?;
        }
    }

    Ok(rewrite.next_catalog_id + allocated)
}

fn rewritten_sort_expression(
    mut expression: SortExpression,
    new_sort_id: SortId,
    old_column_name: &str,
    new_column_name: &str,
) -> Result<SortExpression> {
    expression.sort_id = new_sort_id;
    expression.expression = crate::rewrite_sort_expression_for_renamed_column(
        &expression.expression,
        old_column_name,
        new_column_name,
    )?;
    Ok(expression)
}

fn validate_closed_sort_count(table_id: TableId, sort_id: SortId, closed: usize) -> Result<()> {
    if closed == 1 {
        return Ok(());
    }
    Err(StorageError::Backend(format!(
        "DuckLake RENAME COLUMN expected one active sort row for table {} sort {}, updated {}",
        table_id.0, sort_id.0, closed
    )))
}
