tempest-engine 0.0.2

use std::str::FromStr;

use derive_more::{Display, Error};
use itertools::Itertools;
use tempest_core::tempest_str::TempestStr;
use tempest_tql::ast::{self, Path};

use crate::{
    catalog::{
        CatalogState, flatten_schema, pk_path_to_flat_idx,
        schema::{DatabaseId, DatabaseSchema, TypeExpr, TypeId, TypeSchema},
    },
    row::resolved::ResolvedTable,
    types::TempestType,
};

#[derive(Debug, Display, Error)]
pub enum ResolveError {
    #[display("database with the name '{}' was not found", _0)]
    DatabaseNotFound(#[error(not(source))] TempestStr<'static>),
    #[display("table with the name '{}' was not found inside of this scope", _0)]
    TableNotFound(#[error(not(source))] TempestStr<'static>),
    #[display("unqualified path: please specify scope")]
    UnqualifiedPath,
    #[display("type with the name '{}' was not found inside of this scope", _0)]
    UnknownType(#[error(not(source))] TempestStr<'static>),
}

pub(crate) fn resolve_database<'a>(
    path: &Path,
    catalog: &'a CatalogState,
) -> Result<(DatabaseId, &'a DatabaseSchema), ResolveError> {
    let name = path.database().ok_or(ResolveError::UnqualifiedPath)?;
    catalog
        .get_database_by_name(&name.name)
        .ok_or_else(|| ResolveError::DatabaseNotFound(name.name.clone().into_owned()))
}

pub(crate) fn resolve_type_schema<'a>(
    path: &Path<'a>,
    catalog: &'a CatalogState,
) -> Result<(TypeId, &'a TypeSchema), ResolveError> {
    let (db_id, _) = resolve_database(path, catalog)?;
    let name = path.name().ok_or(ResolveError::UnqualifiedPath)?;
    catalog
        .get_type_by_name(db_id, &name.name)
        .ok_or_else(|| ResolveError::UnknownType(name.name.clone().into_owned()))
}

pub(crate) fn resolve_type(
    ty: &ast::TypeExpr,
    catalog: &CatalogState,
    generic_params: &[&TempestStr<'_>],
) -> Result<TypeExpr, ResolveError> {
    let resolved_args = ty
        .generic_args
        .args
        .iter()
        .map(|arg| resolve_type(arg, catalog, generic_params))
        .try_collect()?;

    if let Some(database) = ty.path.database() {
        let Some((db, _)) = catalog.get_database_by_name(&database.name) else {
            return Err(ResolveError::DatabaseNotFound(
                database.name.clone().into_owned(),
            ));
        };
        let type_name = ty.path.name().ok_or(ResolveError::UnqualifiedPath)?;
        let Some((resolved_ty, _)) = catalog.get_type_by_name(db, &type_name.name) else {
            return Err(ResolveError::UnknownType(
                type_name.name.clone().into_owned(),
            ));
        };
        return Ok(TypeExpr::Ref(resolved_ty, resolved_args));
    } else {
        let type_name = ty.path.name().ok_or(ResolveError::UnqualifiedPath)?;
        if let Some(idx) = generic_params.iter().position(|&p| p == &type_name.name) {
            return Ok(TypeExpr::GenericParam(idx as u32));
        }

        if let Ok(resolved_ty) = TempestType::from_str(&type_name.name) {
            return Ok(TypeExpr::Primitive(resolved_ty));
        }

        if let Some((type_id, _)) = catalog.get_global_type_by_name(&type_name.name) {
            return Ok(TypeExpr::Ref(type_id, resolved_args));
        }
    }

    Err(ResolveError::UnknownType(
        ty.path.name().unwrap().name.clone().into_owned(),
    ))
}

pub(crate) fn resolve_table<'a>(
    path: &Path,
    // when we implement session state / `use`, we may supply a database scope as context
    // scope: Option<DatabaseId>,
    catalog: &'a CatalogState,
) -> Result<ResolvedTable<'a>, ResolveError> {
    let (database_id, _) = resolve_database(path, catalog)?;

    let table_name = &path.name().ok_or(ResolveError::UnqualifiedPath)?.name;
    let (table_id, table_schema) = catalog
        .get_table_by_name(database_id, table_name)
        .ok_or_else(|| ResolveError::TableNotFound(table_name.clone().into_owned()))?;

    let struct_schema = catalog
        .get_type(table_schema.type_id)
        .expect("type referenced by table not found in catalog - catalog is corrupt")
        .as_struct()
        .expect("table type must be a struct");

    let flat_fields = flatten_schema(
        &struct_schema.fields,
        &table_schema.generic_args,
        catalog,
        "",
    )
    .expect("flat schema build failed - catalog is inconsistent");

    let primary_key = table_schema.primary_key.iter()
        .map(|path| pk_path_to_flat_idx(
            path,
            &struct_schema.fields,
            &table_schema.generic_args,
            catalog,
            &flat_fields,
        ).expect("pk path not found in flat fields - catalog is inconsistent"))
        .collect();

    Ok(ResolvedTable {
        id: table_id,
        fields: &struct_schema.fields,
        generic_args: &table_schema.generic_args,
        primary_key,
        flat_fields,
    })
}

#[cfg(test)]
mod tests {
    use crate::catalog::{schema::TableId, testing::create_catalog_state_for_testing};

    use super::*;

    #[test]
    fn resolve_basic() {
        let state = create_catalog_state_for_testing();
        let path = Path::for_testing(Some("main".into()), "users".into());
        let resolved = resolve_table(&path, &state).unwrap();
        assert_eq!(resolved.id, TableId(0));
        assert_eq!(resolved.fields.len(), 2);
    }

    #[test]
    fn resolve_database_not_found() {
        let state = create_catalog_state_for_testing();
        let path = Path::for_testing(Some("missing".into()), "users".into());
        assert!(matches!(
            resolve_table(&path, &state),
            Err(ResolveError::DatabaseNotFound(_))
        ));
    }

    #[test]
    fn resolve_table_not_found() {
        let state = create_catalog_state_for_testing();
        let path = Path::for_testing(Some("main".into()), "missing".into());
        assert!(matches!(
            resolve_table(&path, &state),
            Err(ResolveError::TableNotFound(_))
        ));
    }

    #[test]
    fn resolve_unqualified_path() {
        let state = create_catalog_state_for_testing();
        let path = Path::for_testing(None, "users".into());
        assert!(matches!(
            resolve_table(&path, &state),
            Err(ResolveError::UnqualifiedPath)
        ));
    }
}