//! Resolved table providers keyed by Hamelin [`QualifiedDatasetIdentifier`].
//!
//! The DataFusion executor keeps a persistent [`DatasetTableCatalog`] as its
//! provider index; each query also builds a per-query catalog for translation.
//! Hamelin DataFusion translation looks up providers here instead of using
//! DataFusion's three-part [`TableReference`] session catalog.

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

use datafusion::catalog::TableProvider;
use hamelin_lib::catalog::CatalogProvider;
use hamelin_lib::tree::ast::dataset_identifier::QualifiedDatasetIdentifier;
use hamelin_lib::tree::ast::identifier::SimpleIdentifier;
use hamelin_lib::types::Type;
use ordermap::OrderMap;

use crate::arrow::arrow_schema_to_columns;

/// Table providers keyed by canonical Hamelin dataset identifier.
#[derive(Debug, Default, Clone)]
pub struct DatasetTableCatalog {
    tables: HashMap<QualifiedDatasetIdentifier, Arc<dyn TableProvider>>,
}

impl DatasetTableCatalog {
    pub fn new() -> Self {
        Self {
            tables: HashMap::new(),
        }
    }

    pub fn insert(&mut self, id: QualifiedDatasetIdentifier, provider: Arc<dyn TableProvider>) {
        self.tables.insert(id, provider);
    }

    pub fn get(&self, id: &QualifiedDatasetIdentifier) -> Option<Arc<dyn TableProvider>> {
        self.tables.get(id).cloned()
    }

    pub fn provider(
        &self,
        id: &QualifiedDatasetIdentifier,
    ) -> Result<Arc<dyn TableProvider>, String> {
        self.get(id)
            .ok_or_else(|| format!("dataset '{}' is not registered", id))
    }

    pub fn is_empty(&self) -> bool {
        self.tables.is_empty()
    }

    pub fn iter(
        &self,
    ) -> impl Iterator<Item = (&QualifiedDatasetIdentifier, &Arc<dyn TableProvider>)> {
        self.tables.iter()
    }

    /// Build a [`CatalogProvider`] from registered table schemas for type-checking.
    pub fn to_catalog_provider(&self) -> Result<CatalogProvider, String> {
        let catalog = CatalogProvider::default();
        for (id, provider) in self.iter() {
            let cols_map: OrderMap<SimpleIdentifier, Type> =
                arrow_schema_to_columns(&provider.schema())
                    .into_iter()
                    .map(|col| -> Result<(SimpleIdentifier, Type), String> {
                        let name: SimpleIdentifier = col.name.clone().into();
                        let typ: Type = col
                            .typ
                            .try_into()
                            .map_err(|e| format!("invalid column type for '{}': {e}", col.name))?;
                        Ok((name, typ))
                    })
                    .collect::<Result<_, _>>()?;
            catalog.set(id.clone(), cols_map);
        }
        Ok(catalog)
    }
}
