lakehouse/crates/catalogd/src/registry.rs

use shared::types::{
    AiView, ColumnMeta, DatasetId, DatasetManifest, FreshnessContract, Lineage, ModelProfile,
    ObjectRef, RefreshPolicy, SchemaFingerprint, Sensitivity, Tombstone,
};

use crate::tombstones::TombstoneStore;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

use storaged::ops;
use object_store::ObjectStore;

/// Make a view name safe for use as an object storage key.
/// Allows letters, digits, `_`, `-`, `.`. Anything else becomes `_`.
fn sanitize_view_name(name: &str) -> String {
    name.chars()
        .map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.' { c } else { '_' })
        .collect()
}

#[derive(Debug, Clone, Default, serde::Serialize)]
pub struct MigrateBucketsReport {
    pub refs_examined: usize,
    pub refs_renamed: usize,    // legacy "data"/"local" → "primary"
    pub refs_stamped: usize,    // empty → "primary"
    pub refs_unchanged: usize,  // already canonical
    pub manifests_persisted: usize,
}

/// Partial metadata update — only set fields are applied.
#[derive(Debug, Clone, Default, serde::Deserialize)]
pub struct MetadataUpdate {
    pub description: Option<String>,
    pub owner: Option<String>,
    pub sensitivity: Option<Sensitivity>,
    pub tags: Option<Vec<String>>,
    pub columns: Option<Vec<ColumnMeta>>,
    pub lineage: Option<Lineage>,
    pub freshness: Option<FreshnessContract>,
    pub row_count: Option<u64>,
    // Phase C embedding freshness
    pub embedding_refresh_policy: Option<RefreshPolicy>,
}

const MANIFEST_PREFIX: &str = "_catalog/manifests";
const VIEW_PREFIX: &str = "_catalog/views";
const PROFILE_PREFIX: &str = "_catalog/profiles";

/// In-memory dataset registry backed by manifest persistence in object storage.
/// Also tracks AiViews (Phase D) — safe projections over base datasets.
/// And tombstones (Phase E) — soft-delete markers applied at query time.
#[derive(Clone)]
pub struct Registry {
    datasets: Arc<RwLock<HashMap<DatasetId, DatasetManifest>>>,
    views: Arc<RwLock<HashMap<String, AiView>>>,
    profiles: Arc<RwLock<HashMap<String, ModelProfile>>>,
    tombstones: TombstoneStore,
    store: Arc<dyn ObjectStore>,
}

impl Registry {
    pub fn new(store: Arc<dyn ObjectStore>) -> Self {
        Self {
            datasets: Arc::new(RwLock::new(HashMap::new())),
            views: Arc::new(RwLock::new(HashMap::new())),
            profiles: Arc::new(RwLock::new(HashMap::new())),
            tombstones: TombstoneStore::new(store.clone()),
            store,
        }
    }

    pub fn tombstones(&self) -> &TombstoneStore {
        &self.tombstones
    }

    /// Add a tombstone for a dataset row. Validates that the dataset
    /// exists. Idempotent at the endpoint layer — callers dedup if needed.
    pub async fn add_tombstone(
        &self,
        dataset: &str,
        row_key_column: &str,
        row_key_value: &str,
        actor: &str,
        reason: &str,
    ) -> Result<Tombstone, String> {
        if self.get_by_name(dataset).await.is_none() {
            return Err(format!("dataset not found: {dataset}"));
        }
        let ts = Tombstone {
            dataset: dataset.to_string(),
            row_key_column: row_key_column.to_string(),
            row_key_value: row_key_value.to_string(),
            deleted_at: chrono::Utc::now(),
            actor: actor.to_string(),
            reason: reason.to_string(),
        };
        self.tombstones.append(&ts).await?;
        Ok(ts)
    }

    pub async fn list_tombstones(&self, dataset: &str) -> Result<Vec<Tombstone>, String> {
        self.tombstones.list(dataset).await
    }

    /// Rebuild in-memory index from persisted manifests + views on startup.
    pub async fn rebuild(&self) -> Result<usize, String> {
        let keys = ops::list(&self.store, Some(MANIFEST_PREFIX)).await?;
        let mut datasets = self.datasets.write().await;
        datasets.clear();
        for key in &keys {
            let data = ops::get(&self.store, key).await?;
            let manifest: DatasetManifest =
                serde_json::from_slice(&data).map_err(|e| e.to_string())?;
            datasets.insert(manifest.id.clone(), manifest);
        }
        let count = datasets.len();
        tracing::info!("catalog rebuilt: {count} datasets loaded");

        // Phase D: load AiView definitions alongside manifests.
        let view_keys = ops::list(&self.store, Some(VIEW_PREFIX)).await.unwrap_or_default();
        let mut views = self.views.write().await;
        views.clear();
        for key in &view_keys {
            if !key.ends_with(".json") { continue; }
            let data = match ops::get(&self.store, key).await {
                Ok(d) => d,
                Err(e) => {
                    tracing::warn!("view '{key}': read failed: {e}");
                    continue;
                }
            };
            match serde_json::from_slice::<AiView>(&data) {
                Ok(view) => { views.insert(view.name.clone(), view); }
                Err(e) => tracing::warn!("view '{key}': parse failed: {e}"),
            }
        }
        if !views.is_empty() {
            tracing::info!("catalog: {} views loaded", views.len());
        }

        // Phase 17: load model profiles.
        let profile_keys = ops::list(&self.store, Some(PROFILE_PREFIX)).await.unwrap_or_default();
        let mut profiles = self.profiles.write().await;
        profiles.clear();
        for key in &profile_keys {
            if !key.ends_with(".json") { continue; }
            let data = match ops::get(&self.store, key).await {
                Ok(d) => d,
                Err(e) => { tracing::warn!("profile '{key}': read failed: {e}"); continue; }
            };
            match serde_json::from_slice::<ModelProfile>(&data) {
                Ok(p) => { profiles.insert(p.id.clone(), p); }
                Err(e) => tracing::warn!("profile '{key}': parse failed: {e}"),
            }
        }
        if !profiles.is_empty() {
            tracing::info!("catalog: {} model profiles loaded", profiles.len());
        }

        Ok(count)
    }

    // --- Phase 17: Model profiles ---

    /// Create or replace a model profile. Validates id slug + ensures
    /// every bound_dataset exists (as either a raw dataset or an AiView).
    pub async fn put_profile(&self, profile: ModelProfile) -> Result<ModelProfile, String> {
        if profile.id.is_empty() {
            return Err("profile id is empty".into());
        }
        if !profile.id.chars().all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_') {
            return Err(format!(
                "profile id '{}' must be alphanumeric + '-' or '_' only", profile.id,
            ));
        }
        for binding in &profile.bound_datasets {
            let exists = self.get_by_name(binding).await.is_some()
                || self.get_view(binding).await.is_some();
            if !exists {
                return Err(format!(
                    "bound dataset '{}' not found as dataset or view", binding,
                ));
            }
        }

        let key = format!("{PROFILE_PREFIX}/{}.json", profile.id);
        let json = serde_json::to_vec_pretty(&profile).map_err(|e| e.to_string())?;
        ops::put(&self.store, &key, json.into()).await?;

        let mut profiles = self.profiles.write().await;
        profiles.insert(profile.id.clone(), profile.clone());
        tracing::info!(
            "profile registered: {} -> ollama={} bindings={:?}",
            profile.id, profile.ollama_name, profile.bound_datasets,
        );
        Ok(profile)
    }

    pub async fn get_profile(&self, id: &str) -> Option<ModelProfile> {
        self.profiles.read().await.get(id).cloned()
    }

    pub async fn list_profiles(&self) -> Vec<ModelProfile> {
        self.profiles.read().await.values().cloned().collect()
    }

    pub async fn delete_profile(&self, id: &str) -> Result<(), String> {
        let key = format!("{PROFILE_PREFIX}/{id}.json");
        ops::delete(&self.store, &key).await?;
        self.profiles.write().await.remove(id);
        Ok(())
    }

    /// Register a new dataset. Persists manifest to storage before updating memory.
    pub async fn register(
        &self,
        name: String,
        schema_fingerprint: SchemaFingerprint,
        objects: Vec<ObjectRef>,
    ) -> Result<DatasetManifest, String> {
        let now = chrono::Utc::now();
        let manifest = DatasetManifest {
            id: DatasetId::new(),
            name,
            schema_fingerprint,
            objects,
            created_at: now,
            updated_at: now,
            description: String::new(),
            owner: String::new(),
            sensitivity: None,
            columns: vec![],
            lineage: None,
            freshness: None,
            tags: vec![],
            row_count: None,
            last_embedded_at: None,
            embedding_stale_since: None,
            embedding_refresh_policy: None,
        };

        // Write-ahead: persist before in-memory update
        let manifest_key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
        let json = serde_json::to_vec_pretty(&manifest).map_err(|e| e.to_string())?;
        ops::put(&self.store, &manifest_key, json.into()).await?;

        let mut datasets = self.datasets.write().await;
        datasets.insert(manifest.id.clone(), manifest.clone());

        tracing::info!("registered dataset: {} ({})", manifest.name, manifest.id);
        Ok(manifest)
    }

    /// Update metadata on an existing dataset (owner, description, tags, sensitivity, etc.)
    pub async fn update_metadata(
        &self,
        name: &str,
        updates: MetadataUpdate,
    ) -> Result<DatasetManifest, String> {
        let mut datasets = self.datasets.write().await;
        let manifest = datasets.values_mut()
            .find(|d| d.name == name)
            .ok_or_else(|| format!("dataset not found: {name}"))?;

        if let Some(desc) = updates.description { manifest.description = desc; }
        if let Some(owner) = updates.owner { manifest.owner = owner; }
        if let Some(sens) = updates.sensitivity { manifest.sensitivity = Some(sens); }
        if let Some(tags) = updates.tags { manifest.tags = tags; }
        if let Some(cols) = updates.columns { manifest.columns = cols; }
        if let Some(lineage) = updates.lineage { manifest.lineage = Some(lineage); }
        if let Some(freshness) = updates.freshness { manifest.freshness = Some(freshness); }
        if let Some(count) = updates.row_count { manifest.row_count = Some(count); }
        if let Some(policy) = updates.embedding_refresh_policy { manifest.embedding_refresh_policy = Some(policy); }
        manifest.updated_at = chrono::Utc::now();

        // Persist
        let manifest_key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
        let json = serde_json::to_vec_pretty(manifest).map_err(|e| e.to_string())?;
        ops::put(&self.store, &manifest_key, json.into()).await?;

        let result = manifest.clone();
        Ok(result)
    }

    /// Get a dataset by ID.
    pub async fn get(&self, id: &DatasetId) -> Option<DatasetManifest> {
        let datasets = self.datasets.read().await;
        datasets.get(id).cloned()
    }

    /// Get a dataset by name.
    pub async fn get_by_name(&self, name: &str) -> Option<DatasetManifest> {
        let datasets = self.datasets.read().await;
        datasets.values().find(|d| d.name == name).cloned()
    }

    /// List all datasets.
    pub async fn list(&self) -> Vec<DatasetManifest> {
        let datasets = self.datasets.read().await;
        datasets.values().cloned().collect()
    }

    /// Re-read the parquet footer(s) for a dataset and repopulate `row_count`
    /// and `columns` from reality. Use this to repair manifests whose
    /// metadata was lost (e.g. migrated from a pre-Phase 10 catalog).
    ///
    /// Does NOT touch owner/description/sensitivity/lineage/tags — only
    /// the structural facts that parquet can tell us authoritatively.
    /// The existing `schema_fingerprint` is updated if the recomputed one
    /// differs; a warning is logged so drift is visible.
    pub async fn resync_from_parquet(&self, name: &str) -> Result<DatasetManifest, String> {
        use shared::arrow_helpers::{fingerprint_schema, parquet_to_record_batches};

        // Snapshot the target manifest so we don't hold the write lock during IO.
        let (id, objects, old_fp) = {
            let datasets = self.datasets.read().await;
            let m = datasets
                .values()
                .find(|d| d.name == name)
                .ok_or_else(|| format!("dataset not found: {name}"))?;
            (m.id.clone(), m.objects.clone(), m.schema_fingerprint.clone())
        };

        if objects.is_empty() {
            return Err(format!("dataset '{name}' has no object references to resync from"));
        }

        let mut total_rows: u64 = 0;
        let mut first_schema: Option<arrow::datatypes::SchemaRef> = None;
        for obj in &objects {
            let data = ops::get(&self.store, &obj.key).await
                .map_err(|e| format!("read {}: {e}", obj.key))?;
            let (schema, batches) = parquet_to_record_batches(&data)
                .map_err(|e| format!("parse {}: {e}", obj.key))?;
            let rows: u64 = batches.iter().map(|b| b.num_rows() as u64).sum();
            total_rows += rows;
            if first_schema.is_none() {
                first_schema = Some(schema);
            }
        }

        let schema = first_schema.ok_or("no schema recovered")?;
        let new_fp = fingerprint_schema(&schema);
        if new_fp != old_fp {
            tracing::warn!(
                "dataset '{}' schema fingerprint drift: {} -> {} (updating to match parquet reality)",
                name, old_fp.0, new_fp.0,
            );
        }

        let columns: Vec<ColumnMeta> = schema
            .fields()
            .iter()
            .map(|f| ColumnMeta {
                name: f.name().clone(),
                data_type: f.data_type().to_string(),
                sensitivity: None,
                description: String::new(),
                is_pii: false,
            })
            .collect();

        // Apply updates.
        let mut datasets = self.datasets.write().await;
        let manifest = datasets
            .get_mut(&id)
            .ok_or_else(|| format!("dataset disappeared during resync: {name}"))?;
        manifest.row_count = Some(total_rows);
        manifest.columns = columns;
        manifest.schema_fingerprint = new_fp;
        manifest.updated_at = chrono::Utc::now();

        // Persist.
        let manifest_key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
        let json = serde_json::to_vec_pretty(manifest).map_err(|e| e.to_string())?;
        ops::put(&self.store, &manifest_key, json.into()).await?;

        tracing::info!("resynced '{name}': row_count={total_rows}, {} columns", manifest.columns.len());
        Ok(manifest.clone())
    }

    /// Resync every dataset that currently has a null row_count.
    /// Returns (successes, failures) where each entry is (name, detail).
    pub async fn resync_missing(&self) -> (Vec<(String, u64)>, Vec<(String, String)>) {
        let names: Vec<String> = {
            let datasets = self.datasets.read().await;
            datasets
                .values()
                .filter(|d| d.row_count.is_none() || d.columns.is_empty())
                .map(|d| d.name.clone())
                .collect()
        };

        let mut ok = Vec::new();
        let mut err = Vec::new();
        for name in names {
            match self.resync_from_parquet(&name).await {
                Ok(m) => ok.push((name, m.row_count.unwrap_or(0))),
                Err(e) => err.push((name, e)),
            }
        }
        (ok, err)
    }

    /// Mark a dataset's embeddings as stale (row-level data has been written
    /// since the last embedding refresh). Idempotent — setting stale when
    /// already stale is a no-op. Only marks stale if the dataset has been
    /// embedded before — a never-embedded dataset doesn't need a stale flag
    /// (it just needs an initial index build). Called from the ingest path.
    pub async fn mark_embeddings_stale(&self, name: &str) -> Result<(), String> {
        let mut datasets = self.datasets.write().await;
        let manifest = datasets
            .values_mut()
            .find(|d| d.name == name)
            .ok_or_else(|| format!("dataset not found: {name}"))?;

        if manifest.last_embedded_at.is_none() {
            return Ok(()); // never embedded -> no stale semantics yet
        }
        if manifest.embedding_stale_since.is_none() {
            manifest.embedding_stale_since = Some(chrono::Utc::now());
            manifest.updated_at = chrono::Utc::now();

            let key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
            let json = serde_json::to_vec_pretty(manifest).map_err(|e| e.to_string())?;
            ops::put(&self.store, &key, json.into()).await?;
            tracing::info!("marked embeddings stale for dataset '{name}'");
        }
        Ok(())
    }

    /// Clear the stale marker and set `last_embedded_at = now`.
    /// Called by the embedding refresh pipeline once it finishes.
    pub async fn clear_embeddings_stale(&self, name: &str) -> Result<(), String> {
        let mut datasets = self.datasets.write().await;
        let manifest = datasets
            .values_mut()
            .find(|d| d.name == name)
            .ok_or_else(|| format!("dataset not found: {name}"))?;

        let now = chrono::Utc::now();
        manifest.embedding_stale_since = None;
        manifest.last_embedded_at = Some(now);
        manifest.updated_at = now;

        let key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
        let json = serde_json::to_vec_pretty(manifest).map_err(|e| e.to_string())?;
        ops::put(&self.store, &key, json.into()).await?;
        Ok(())
    }

    /// Federation layer 2: stamp `bucket = "primary"` on every ObjectRef
    /// whose `bucket` field is empty or matches a legacy value (`"data"`,
    /// `"local"`). One-shot migration; re-running is a safe no-op once
    /// every ref is canonical.
    pub async fn migrate_buckets_to_primary(&self) -> Result<MigrateBucketsReport, String> {
        let mut datasets = self.datasets.write().await;
        let mut report = MigrateBucketsReport::default();
        let mut to_persist: Vec<(DatasetId, DatasetManifest)> = Vec::new();

        for manifest in datasets.values_mut() {
            let mut changed = false;
            for obj in manifest.objects.iter_mut() {
                report.refs_examined += 1;
                if obj.bucket.is_empty() {
                    obj.bucket = "primary".to_string();
                    report.refs_stamped += 1;
                    changed = true;
                } else if obj.bucket == "data" || obj.bucket == "local" {
                    obj.bucket = "primary".to_string();
                    report.refs_renamed += 1;
                    changed = true;
                } else {
                    report.refs_unchanged += 1;
                }
            }
            if changed {
                manifest.updated_at = chrono::Utc::now();
                to_persist.push((manifest.id.clone(), manifest.clone()));
            }
        }

        // Persist updated manifests after we've finished mutating the map.
        for (id, manifest) in to_persist {
            let key = format!("{MANIFEST_PREFIX}/{}.json", id);
            let json = serde_json::to_vec_pretty(&manifest).map_err(|e| e.to_string())?;
            ops::put(&self.store, &key, json.into()).await?;
            report.manifests_persisted += 1;
        }

        tracing::info!(
            "bucket migration: examined {} refs, renamed {}, stamped {}, unchanged {}, persisted {} manifests",
            report.refs_examined,
            report.refs_renamed,
            report.refs_stamped,
            report.refs_unchanged,
            report.manifests_persisted,
        );
        Ok(report)
    }

    // --- Phase D: AI-safe views ---

    /// Create or replace a named view. Validates that the base dataset
    /// exists and that the column whitelist is non-empty. Persists to
    /// `_catalog/views/{name}.json`.
    pub async fn put_view(&self, mut view: AiView) -> Result<AiView, String> {
        if view.name.is_empty() {
            return Err("view name is empty".into());
        }
        if view.columns.is_empty() {
            return Err("view must whitelist at least one column".into());
        }
        // Base dataset must exist (read-side will fail anyway, fail-fast here
        // so operators don't end up with dangling views).
        if self.get_by_name(&view.base_dataset).await.is_none() {
            return Err(format!(
                "base dataset '{}' not found in catalog",
                view.base_dataset
            ));
        }
        if view.created_at.timestamp() == 0 {
            view.created_at = chrono::Utc::now();
        }

        let key = format!("{VIEW_PREFIX}/{}.json", sanitize_view_name(&view.name));
        let json = serde_json::to_vec_pretty(&view).map_err(|e| e.to_string())?;
        ops::put(&self.store, &key, json.into()).await?;

        let mut views = self.views.write().await;
        views.insert(view.name.clone(), view.clone());
        tracing::info!("view registered: {} over '{}'", view.name, view.base_dataset);
        Ok(view)
    }

    pub async fn get_view(&self, name: &str) -> Option<AiView> {
        self.views.read().await.get(name).cloned()
    }

    pub async fn list_views(&self) -> Vec<AiView> {
        self.views.read().await.values().cloned().collect()
    }

    pub async fn delete_view(&self, name: &str) -> Result<(), String> {
        let key = format!("{VIEW_PREFIX}/{}.json", sanitize_view_name(name));
        ops::delete(&self.store, &key).await?;
        self.views.write().await.remove(name);
        Ok(())
    }

    /// List datasets whose `embedding_stale_since` is set — they need a refresh.
    pub async fn stale_datasets(&self) -> Vec<DatasetManifest> {
        let datasets = self.datasets.read().await;
        datasets
            .values()
            .filter(|d| d.embedding_stale_since.is_some())
            .cloned()
            .collect()
    }

    /// Add objects to an existing dataset.
    pub async fn add_objects(
        &self,
        id: &DatasetId,
        new_objects: Vec<ObjectRef>,
    ) -> Result<DatasetManifest, String> {
        let mut datasets = self.datasets.write().await;
        let manifest = datasets.get_mut(id).ok_or_else(|| format!("dataset not found: {id}"))?;

        manifest.objects.extend(new_objects);
        manifest.updated_at = chrono::Utc::now();

        // Persist updated manifest
        let manifest_key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
        let json = serde_json::to_vec_pretty(manifest).map_err(|e| e.to_string())?;
        ops::put(&self.store, &manifest_key, json.into()).await?;

        Ok(manifest.clone())
    }
}