use std::collections::HashMap;

#[cfg(any(feature = "postgres", feature = "quack"))]
use std::future::Future;

use ducklake_catalog::{
    ColumnId, DuckLakeType, FileColumnStats, FileVariantStats, SnapshotId, TableColumnStats,
    TableId, merge_table_column_stats_for_append,
};

use crate::{AppendFile, Result};

#[cfg(any(feature = "sqlite", feature = "duckdb", test))]
pub(crate) fn update_table_column_stats_for_append_sync<
    Exists,
    QueryColumnTypes,
    QueryTableStats,
    Upsert,
>(
    mut table_column_stats_exists: Exists,
    mut query_column_types: QueryColumnTypes,
    mut query_table_column_stats: QueryTableStats,
    mut upsert_table_column_stats: Upsert,
    table_id: TableId,
    snapshot: SnapshotId,
    files: &[AppendFile],
) -> Result<()>
where
    Exists: FnMut() -> Result<bool>,
    QueryColumnTypes:
        FnMut(crate::catalog_queries::CatalogQuery) -> Result<Vec<(ColumnId, DuckLakeType)>>,
    QueryTableStats: FnMut(crate::catalog_queries::CatalogQuery) -> Result<Vec<TableColumnStats>>,
    Upsert: FnMut(&TableColumnStats) -> Result<()>,
{
    if !append_stats_update_needed(files) || !table_column_stats_exists()? {
        return Ok(());
    }

    let column_types = crate::catalog_queries::column_types_by_id(query_column_types(
        crate::catalog_queries::column_types_at(table_id, snapshot),
    )?);
    let existing_stats = crate::catalog_queries::table_column_stats_by_column(
        query_table_column_stats(crate::catalog_queries::table_column_stats(table_id))?,
    );
    for stats in table_column_stats_for_append(table_id, &column_types, &existing_stats, files) {
        upsert_table_column_stats(&stats)?;
    }
    Ok(())
}

#[cfg(any(feature = "postgres", feature = "quack"))]
pub(crate) async fn update_table_column_stats_for_append_async<
    Exists,
    ExistsFuture,
    QueryColumnTypes,
    QueryColumnTypesFuture,
    QueryTableStats,
    QueryTableStatsFuture,
    Upsert,
    UpsertFuture,
>(
    mut table_column_stats_exists: Exists,
    mut query_column_types: QueryColumnTypes,
    mut query_table_column_stats: QueryTableStats,
    mut upsert_table_column_stats: Upsert,
    table_id: TableId,
    snapshot: SnapshotId,
    files: &[AppendFile],
) -> Result<()>
where
    Exists: FnMut() -> ExistsFuture,
    ExistsFuture: Future<Output = Result<bool>>,
    QueryColumnTypes: FnMut(crate::catalog_queries::CatalogQuery) -> QueryColumnTypesFuture,
    QueryColumnTypesFuture: Future<Output = Result<Vec<(ColumnId, DuckLakeType)>>>,
    QueryTableStats: FnMut(crate::catalog_queries::CatalogQuery) -> QueryTableStatsFuture,
    QueryTableStatsFuture: Future<Output = Result<Vec<TableColumnStats>>>,
    Upsert: FnMut(TableColumnStats) -> UpsertFuture,
    UpsertFuture: Future<Output = Result<()>>,
{
    if !append_stats_update_needed(files) || !table_column_stats_exists().await? {
        return Ok(());
    }

    let column_types = crate::catalog_queries::column_types_by_id(
        query_column_types(crate::catalog_queries::column_types_at(table_id, snapshot)).await?,
    );
    let existing_stats = crate::catalog_queries::table_column_stats_by_column(
        query_table_column_stats(crate::catalog_queries::table_column_stats(table_id)).await?,
    );
    for stats in table_column_stats_for_append(table_id, &column_types, &existing_stats, files) {
        upsert_table_column_stats(stats).await?;
    }
    Ok(())
}

#[cfg(any(feature = "sqlite", test))]
pub(crate) fn insert_table_stats_from_active_files_sync<QueryActiveStats, Execute>(
    table_id: TableId,
    mut query_active_stats: QueryActiveStats,
    mut execute: Execute,
) -> Result<()>
where
    QueryActiveStats: FnMut(
        crate::catalog_queries::CatalogQuery,
    ) -> Result<crate::catalog_queries::ActiveTableStats>,
    Execute: FnMut(crate::catalog_queries::CatalogQuery) -> Result<usize>,
{
    let stats = query_active_stats(crate::catalog_queries::active_table_stats_from_data_files(
        table_id,
    ))?;
    execute(crate::catalog_queries::insert_table_stats(
        table_id,
        stats.record_count,
        stats.next_row_id,
        stats.file_size_bytes,
    ))?;
    Ok(())
}

#[cfg(any(feature = "sqlite", test))]
pub(crate) fn recompute_table_column_stats_from_active_files_sync<
    Exists,
    Execute,
    QueryColumnTypes,
    QueryActiveFileCount,
    QueryFileStats,
    QueryVariantStats,
    Upsert,
>(
    table_id: TableId,
    snapshot: SnapshotId,
    mut table_exists: Exists,
    mut execute: Execute,
    mut query_column_types: QueryColumnTypes,
    mut query_active_file_count: QueryActiveFileCount,
    mut query_file_stats: QueryFileStats,
    mut query_variant_stats: QueryVariantStats,
    mut upsert_table_column_stats: Upsert,
) -> Result<()>
where
    Exists: FnMut(&str) -> Result<bool>,
    Execute: FnMut(crate::catalog_queries::CatalogQuery) -> Result<usize>,
    QueryColumnTypes:
        FnMut(crate::catalog_queries::CatalogQuery) -> Result<Vec<(ColumnId, DuckLakeType)>>,
    QueryActiveFileCount: FnMut(crate::catalog_queries::CatalogQuery) -> Result<i64>,
    QueryFileStats: FnMut(crate::catalog_queries::CatalogQuery) -> Result<Vec<FileColumnStats>>,
    QueryVariantStats: FnMut(crate::catalog_queries::CatalogQuery) -> Result<Vec<FileVariantStats>>,
    Upsert: FnMut(&TableColumnStats) -> Result<()>,
{
    if !table_exists("ducklake_table_column_stats")? {
        return Ok(());
    }
    let column_types = crate::catalog_queries::column_types_by_id(query_column_types(
        crate::catalog_queries::column_types_at(table_id, snapshot),
    )?);
    execute(crate::catalog_queries::delete_table_column_stats(table_id))?;

    let active_file_count = active_data_file_count(query_active_file_count(
        crate::catalog_queries::active_data_file_count(table_id),
    )?)?;
    if active_file_count == 0 {
        return Ok(());
    }

    let active_by_column = if table_exists("ducklake_file_column_stats")? {
        file_column_stats_by_column(query_file_stats(
            crate::catalog_queries::active_file_column_stats(table_id),
        )?)
    } else {
        HashMap::new()
    };
    let active_variant_by_column = if table_exists("ducklake_file_variant_stats")? {
        file_variant_stats_by_column(query_variant_stats(
            crate::catalog_queries::active_file_variant_stats(table_id),
        )?)
    } else {
        HashMap::new()
    };
    for stats in table_column_stats_from_active_files(
        table_id,
        &column_types,
        &active_by_column,
        &active_variant_by_column,
        active_file_count,
    ) {
        upsert_table_column_stats(&stats)?;
    }
    Ok(())
}

#[cfg(any(feature = "sqlite", test))]
fn active_data_file_count(count: i64) -> Result<usize> {
    usize::try_from(count).map_err(|_| {
        crate::StorageError::Backend("DuckLake active data file count is out of range".into())
    })
}

#[cfg(any(feature = "sqlite", test))]
fn file_column_stats_by_column(
    rows: impl IntoIterator<Item = FileColumnStats>,
) -> HashMap<ColumnId, Vec<FileColumnStats>> {
    let mut by_column: HashMap<ColumnId, Vec<FileColumnStats>> = HashMap::new();
    for stats in rows {
        by_column.entry(stats.column_id).or_default().push(stats);
    }
    by_column
}

#[cfg(any(feature = "sqlite", test))]
fn file_variant_stats_by_column(
    rows: impl IntoIterator<Item = FileVariantStats>,
) -> HashMap<ColumnId, Vec<FileVariantStats>> {
    let mut by_column: HashMap<ColumnId, Vec<FileVariantStats>> = HashMap::new();
    for stats in rows {
        by_column.entry(stats.column_id).or_default().push(stats);
    }
    by_column
}

#[cfg(any(feature = "duckdb", test))]
pub(crate) fn next_row_id_for_append_from_table_stats_sync_query<RunQuery>(
    mut query_optional_i64: RunQuery,
    table_id: TableId,
    has_table_stats: bool,
) -> Result<Option<i64>>
where
    RunQuery: FnMut(crate::catalog_queries::CatalogQuery) -> Result<Option<i64>>,
{
    if !has_table_stats {
        return Ok(None);
    }
    query_optional_i64(crate::catalog_queries::coalesced_next_row_id_from_table_stats(table_id))
}

#[cfg(any(feature = "sqlite", test))]
pub(crate) fn next_row_id_for_append_with_data_file_fallback_sync_query<RunQuery>(
    mut query_optional_i64: RunQuery,
    table_id: TableId,
    has_table_stats: bool,
) -> Result<Option<i64>>
where
    RunQuery: FnMut(crate::catalog_queries::CatalogQuery) -> Result<Option<i64>>,
{
    let stats_next_row_id = if has_table_stats {
        query_optional_i64(crate::catalog_queries::next_row_id_from_table_stats(
            table_id,
        ))?
    } else {
        None
    };
    let data_files_next_row_id = query_optional_i64(
        crate::catalog_queries::next_row_id_from_data_files(table_id),
    )?;
    Ok(stats_next_row_id
        .or(data_files_next_row_id)
        .or_else(|| has_table_stats.then_some(0)))
}

#[cfg(any(feature = "postgres", feature = "quack"))]
pub(crate) async fn next_row_id_for_append_from_table_stats_async_query<RunQuery, QueryFuture>(
    mut query_optional_i64: RunQuery,
    table_id: TableId,
    has_table_stats: bool,
) -> Result<Option<i64>>
where
    RunQuery: FnMut(crate::catalog_queries::CatalogQuery) -> QueryFuture,
    QueryFuture: Future<Output = Result<Option<i64>>>,
{
    if !has_table_stats {
        return Ok(None);
    }
    query_optional_i64(crate::catalog_queries::coalesced_next_row_id_from_table_stats(table_id))
        .await
}

#[cfg(any(feature = "sqlite", feature = "duckdb", test))]
pub(crate) fn upsert_table_column_stats_by_update_sync<Execute>(
    stats: &TableColumnStats,
    mut execute: Execute,
) -> Result<()>
where
    Execute: FnMut(crate::catalog_queries::CatalogQuery) -> Result<usize>,
{
    let updated = execute(crate::catalog_queries::update_table_column_stats(stats))?;
    if updated == 0 {
        execute(crate::catalog_queries::insert_table_column_stats(stats))?;
    }
    Ok(())
}

#[cfg(feature = "postgres")]
pub(crate) async fn upsert_table_column_stats_by_update_async<Execute, ExecuteFuture>(
    stats: &TableColumnStats,
    mut execute: Execute,
) -> Result<()>
where
    Execute: FnMut(crate::catalog_queries::CatalogQuery) -> ExecuteFuture,
    ExecuteFuture: Future<Output = Result<usize>>,
{
    let updated = execute(crate::catalog_queries::update_table_column_stats(stats)).await?;
    if updated == 0 {
        execute(crate::catalog_queries::insert_table_column_stats(stats)).await?;
    }
    Ok(())
}

#[cfg(feature = "quack")]
pub(crate) async fn upsert_table_column_stats_by_count_async<
    Count,
    CountFuture,
    Execute,
    ExecuteFuture,
>(
    stats: &TableColumnStats,
    count: Count,
    execute: Execute,
) -> Result<()>
where
    Count: FnOnce(crate::catalog_queries::CatalogQuery) -> CountFuture,
    CountFuture: Future<Output = Result<i64>>,
    Execute: FnOnce(crate::catalog_queries::CatalogQuery) -> ExecuteFuture,
    ExecuteFuture: Future<Output = Result<()>>,
{
    let existing = count(crate::catalog_queries::table_column_stats_count(
        stats.table_id,
        stats.column_id,
    ))
    .await?;
    let query = if existing > 0 {
        crate::catalog_queries::update_table_column_stats(stats)
    } else {
        crate::catalog_queries::insert_table_column_stats(stats)
    };
    execute(query).await
}

fn append_stats_update_needed(files: &[AppendFile]) -> bool {
    !files.is_empty()
        && files
            .iter()
            .any(|file| !file.column_stats.is_empty() || !file.variant_stats.is_empty())
}

pub(crate) fn table_column_stats_for_append(
    table_id: TableId,
    column_types: &HashMap<ColumnId, DuckLakeType>,
    existing_stats: &HashMap<ColumnId, TableColumnStats>,
    files: &[AppendFile],
) -> Vec<TableColumnStats> {
    let mut appended_by_column: HashMap<ColumnId, Vec<&FileColumnStats>> = HashMap::new();
    let mut appended_variant_by_column: HashMap<ColumnId, Vec<&FileVariantStats>> = HashMap::new();
    for file in files {
        for stats in &file.column_stats {
            appended_by_column
                .entry(stats.column_id)
                .or_default()
                .push(stats);
        }
        for stats in &file.variant_stats {
            appended_variant_by_column
                .entry(stats.column_id)
                .or_default()
                .push(stats);
        }
    }

    merged_table_column_stats(
        table_id,
        column_types,
        existing_stats,
        &appended_by_column,
        &appended_variant_by_column,
        files.len(),
    )
}

#[cfg(any(feature = "sqlite", test))]
pub(crate) fn table_column_stats_from_active_files(
    table_id: TableId,
    column_types: &HashMap<ColumnId, DuckLakeType>,
    active_by_column: &HashMap<ColumnId, Vec<FileColumnStats>>,
    active_variant_by_column: &HashMap<ColumnId, Vec<FileVariantStats>>,
    active_file_count: usize,
) -> Vec<TableColumnStats> {
    let active_refs = active_by_column
        .iter()
        .map(|(column_id, stats)| (*column_id, stats.iter().collect::<Vec<_>>()))
        .collect::<HashMap<_, _>>();
    let active_variant_refs = active_variant_by_column
        .iter()
        .map(|(column_id, stats)| (*column_id, stats.iter().collect::<Vec<_>>()))
        .collect::<HashMap<_, _>>();

    merged_table_column_stats(
        table_id,
        column_types,
        &HashMap::new(),
        &active_refs,
        &active_variant_refs,
        active_file_count,
    )
}

fn merged_table_column_stats(
    table_id: TableId,
    column_types: &HashMap<ColumnId, DuckLakeType>,
    existing_stats: &HashMap<ColumnId, TableColumnStats>,
    appended_by_column: &HashMap<ColumnId, Vec<&FileColumnStats>>,
    appended_variant_by_column: &HashMap<ColumnId, Vec<&FileVariantStats>>,
    appended_file_count: usize,
) -> Vec<TableColumnStats> {
    let mut column_ids = appended_by_column
        .keys()
        .chain(appended_variant_by_column.keys())
        .copied()
        .collect::<Vec<_>>();
    column_ids.sort();
    column_ids.dedup();

    column_ids
        .into_iter()
        .filter_map(|column_id| {
            let column_type = column_types.get(&column_id)?;
            let appended = appended_by_column
                .get(&column_id)
                .map(Vec::as_slice)
                .unwrap_or(&[]);
            let appended_variant = appended_variant_by_column
                .get(&column_id)
                .map(Vec::as_slice)
                .unwrap_or(&[]);
            Some(merge_table_column_stats_for_append(
                existing_stats.get(&column_id),
                table_id,
                column_id,
                column_type,
                appended,
                appended_variant,
                appended_file_count,
            ))
        })
        .collect()
}
