NxMesh/apps/nxmesh-master/src/service/proxy/repo.rs

use std::collections::HashMap;

use sea_orm::{DatabaseConnection, prelude::*};

use crate::service::proxy::types::{
    Mergeable, OverrideRef, ProxyConfig, ProxyServiceError, ProxyServiceResult, ProxyType,
};

#[async_trait::async_trait]
pub trait ProxyRepo: Send + Sync + 'static {
    // get the raw config for the given proxy_id. This should return the config of the given proxy_id without merging it with its parent configs (if any).
    async fn get_proxy_raw_config(&self, proxy_id: uuid::Uuid) -> ProxyServiceResult<ProxyConfig>;

    // get the raw config for the given proxy_id. This should return the config of the given proxy_id and all its parent configs (if any) without merging them. The returned vector is ordered from the leaf (given proxy_id) down to the root config (most specific to least specific).
    async fn get_proxy_raw_configs(
        &self,
        proxy_id: uuid::Uuid,
    ) -> ProxyServiceResult<Vec<ProxyConfig>>;

    // get the merged config for the given proxy_id. This should merge the config of the given proxy_id with its parent configs (if any) and return the final merged config.
    async fn get_merged_proxy_config(
        &self,
        proxy_id: uuid::Uuid,
    ) -> ProxyServiceResult<ProxyConfig>;
}

pub(crate) struct ProxyRepoImpl {
    db: DatabaseConnection,
}

impl ProxyRepoImpl {
    pub fn new(db: DatabaseConnection) -> Self {
        Self { db }
    }
}

#[async_trait::async_trait]
impl ProxyRepo for ProxyRepoImpl {
    async fn get_proxy_raw_configs(
        &self,
        proxy_config_id: uuid::Uuid,
    ) -> ProxyServiceResult<Vec<ProxyConfig>> {
        let mut proxy_config_id_frontier = vec![proxy_config_id];
        let mut visited = std::collections::HashSet::new();

        let mut configs: Vec<ProxyConfig> = Vec::new();

        while let Some(current_id) = proxy_config_id_frontier.pop() {
            if visited.contains(&current_id) {
                continue;
            }
            visited.insert(current_id);
            //
            let config = self.get_proxy_raw_config(current_id).await?;
            //
            if let Some(parent_ids) = &config.parent_config_id {
                proxy_config_id_frontier.extend(parent_ids);
            }
            configs.push(config);
        }

        Ok(configs)
    }

    async fn get_proxy_raw_config(&self, proxy_id: uuid::Uuid) -> ProxyServiceResult<ProxyConfig> {
        let proxy_entity = crate::db::entities::proxy_config::Entity::find_by_id(proxy_id)
            .one(&self.db)
            .await?
            .ok_or(ProxyServiceError::ConfigNotFound)?;

        // Resolve parent configs from config_inheritance
        let parent_config_id = {
            let inheritance = crate::db::entities::config_inheritance::Entity::find()
                .filter(crate::db::entities::config_inheritance::Column::ChildConfigId.eq(proxy_id))
                .all(&self.db)
                .await?;
            if inheritance.is_empty() {
                None
            } else {
                Some(
                    inheritance
                        .into_iter()
                        .map(|ci| ci.parent_config_id)
                        .collect(),
                )
            }
        };

        // Direct children of proxy_config
        let server_blocks = crate::db::entities::server_block::Entity::find()
            .filter(crate::db::entities::server_block::Column::ConfigId.eq(proxy_id))
            .all(&self.db)
            .await?;

        let upstreams = crate::db::entities::upstream::Entity::find()
            .filter(crate::db::entities::upstream::Column::ConfigId.eq(proxy_id))
            .all(&self.db)
            .await?;

        // Children of server_blocks
        let server_block_ids: Vec<uuid::Uuid> = server_blocks.iter().map(|sb| sb.id).collect();

        let (location_blocks, log_settings, server_access_rules) = if server_block_ids.is_empty() {
            (vec![], vec![], vec![])
        } else {
            (
                crate::db::entities::location_block::Entity::find()
                    .filter(
                        crate::db::entities::location_block::Column::ServerId
                            .is_in(server_block_ids.clone()),
                    )
                    .all(&self.db)
                    .await?,
                crate::db::entities::log_setting::Entity::find()
                    .filter(
                        crate::db::entities::log_setting::Column::ServerId
                            .is_in(server_block_ids.clone()),
                    )
                    .all(&self.db)
                    .await?,
                crate::db::entities::access_rule::Entity::find()
                    .filter(
                        crate::db::entities::access_rule::Column::ServerId.is_in(server_block_ids),
                    )
                    .all(&self.db)
                    .await?,
            )
        };

        // Children of location_blocks
        let location_block_ids: Vec<uuid::Uuid> = location_blocks.iter().map(|lb| lb.id).collect();

        let (location_access_rules, limit_rules, proxy_settings, rewrite_rules) =
            if location_block_ids.is_empty() {
                (vec![], vec![], vec![], vec![])
            } else {
                (
                    crate::db::entities::access_rule::Entity::find()
                        .filter(
                            crate::db::entities::access_rule::Column::LocationId
                                .is_in(location_block_ids.clone()),
                        )
                        .all(&self.db)
                        .await?,
                    crate::db::entities::limit_rule::Entity::find()
                        .filter(
                            crate::db::entities::limit_rule::Column::LocationId
                                .is_in(location_block_ids.clone()),
                        )
                        .all(&self.db)
                        .await?,
                    crate::db::entities::proxy_setting::Entity::find()
                        .filter(
                            crate::db::entities::proxy_setting::Column::LocationId
                                .is_in(location_block_ids.clone()),
                        )
                        .all(&self.db)
                        .await?,
                    crate::db::entities::rewrite_rule::Entity::find()
                        .filter(
                            crate::db::entities::rewrite_rule::Column::LocationId
                                .is_in(location_block_ids),
                        )
                        .all(&self.db)
                        .await?,
                )
            };

        // Collect referenced IDs for zone and cert lookups
        let limit_zone_ids: Vec<uuid::Uuid> = limit_rules.iter().map(|lr| lr.zone_id).collect();
        let cache_zone_ids: Vec<uuid::Uuid> = proxy_settings
            .iter()
            .filter_map(|ps| ps.cache_zone)
            .collect();
        let ssl_cert_ids: Vec<uuid::Uuid> = server_blocks
            .iter()
            .filter_map(|sb| sb.ssl_cert_id)
            .collect();

        let limit_zones = if limit_zone_ids.is_empty() {
            vec![]
        } else {
            crate::db::entities::limit_zone::Entity::find()
                .filter(crate::db::entities::limit_zone::Column::Id.is_in(limit_zone_ids))
                .all(&self.db)
                .await?
        };

        let cache_zones = if cache_zone_ids.is_empty() {
            vec![]
        } else {
            crate::db::entities::cache_zone::Entity::find()
                .filter(crate::db::entities::cache_zone::Column::Id.is_in(cache_zone_ids))
                .all(&self.db)
                .await?
        };

        let ssl_certificates = if ssl_cert_ids.is_empty() {
            vec![]
        } else {
            crate::db::entities::ssl_certificate::Entity::find()
                .filter(crate::db::entities::ssl_certificate::Column::Id.is_in(ssl_cert_ids))
                .all(&self.db)
                .await?
        };

        // ── Group child IDs by parent ────────────────────────────────────────

        // location_block IDs grouped by server_id
        let loc_block_ids_by_server: HashMap<uuid::Uuid, Vec<OverrideRef>> = {
            let mut map: HashMap<_, Vec<_>> = HashMap::new();
            for lb in &location_blocks {
                map.entry(lb.server_id).or_default().push(OverrideRef {
                    id: lb.id,
                    override_of_id: lb.override_of_id,
                });
            }
            map
        };

        // log_setting IDs grouped by server_id
        let log_setting_ids_by_server: HashMap<uuid::Uuid, Vec<OverrideRef>> = {
            let mut map: HashMap<_, Vec<_>> = HashMap::new();
            for ls in &log_settings {
                map.entry(ls.server_id).or_default().push(OverrideRef {
                    id: ls.id,
                    override_of_id: ls.override_of_id,
                });
            }
            map
        };

        // server-level access_rule IDs grouped by server_id
        let server_ar_ids_by_server: HashMap<uuid::Uuid, Vec<OverrideRef>> = {
            let mut map: HashMap<_, Vec<_>> = HashMap::new();
            for ar in &server_access_rules {
                if let Some(sid) = ar.server_id {
                    map.entry(sid).or_default().push(OverrideRef {
                        id: ar.id,
                        override_of_id: ar.override_of_id,
                    });
                }
            }
            map
        };

        // location-level access_rule IDs grouped by location_id
        let loc_ar_ids_by_location: HashMap<uuid::Uuid, Vec<OverrideRef>> = {
            let mut map: HashMap<_, Vec<_>> = HashMap::new();
            for ar in &location_access_rules {
                if let Some(lid) = ar.location_id {
                    map.entry(lid).or_default().push(OverrideRef {
                        id: ar.id,
                        override_of_id: ar.override_of_id,
                    });
                }
            }
            map
        };

        // limit_rule IDs grouped by location_id
        let lr_ids_by_location: HashMap<uuid::Uuid, Vec<OverrideRef>> = {
            let mut map: HashMap<_, Vec<_>> = HashMap::new();
            for lr in &limit_rules {
                map.entry(lr.location_id).or_default().push(OverrideRef {
                    id: lr.id,
                    override_of_id: lr.override_of_id,
                });
            }
            map
        };

        // proxy_setting IDs grouped by location_id
        let ps_ids_by_location: HashMap<uuid::Uuid, Vec<OverrideRef>> = {
            let mut map: HashMap<_, Vec<_>> = HashMap::new();
            for ps in &proxy_settings {
                map.entry(ps.location_id).or_default().push(OverrideRef {
                    id: ps.id,
                    override_of_id: ps.override_of_id,
                });
            }
            map
        };

        // rewrite_rule IDs grouped by location_id
        let rr_ids_by_location: HashMap<uuid::Uuid, Vec<OverrideRef>> = {
            let mut map: HashMap<_, Vec<_>> = HashMap::new();
            for rr in &rewrite_rules {
                map.entry(rr.location_id).or_default().push(OverrideRef {
                    id: rr.id,
                    override_of_id: rr.override_of_id,
                });
            }
            map
        };

        // location_block IDs grouped by proxy_pass_upstream_id (reverse FK)
        let loc_block_ids_by_upstream: HashMap<uuid::Uuid, Vec<OverrideRef>> = {
            let mut map: HashMap<_, Vec<_>> = HashMap::new();
            for lb in &location_blocks {
                if let Some(up_id) = lb.proxy_pass_upstream_id {
                    map.entry(up_id).or_default().push(OverrideRef {
                        id: lb.id,
                        override_of_id: lb.override_of_id,
                    });
                }
            }
            map
        };

        // ── Build child-ID tuples for each config type ───────────────────────

        let location_block_tuples: Vec<(
            crate::db::entities::location_block::Model,
            Vec<OverrideRef>,
            Vec<OverrideRef>,
            Vec<OverrideRef>,
            Vec<OverrideRef>,
        )> = location_blocks
            .into_iter()
            .map(|lb| {
                let ar_ids = loc_ar_ids_by_location
                    .get(&lb.id)
                    .cloned()
                    .unwrap_or_default();
                let lr_ids = lr_ids_by_location.get(&lb.id).cloned().unwrap_or_default();
                let ps_ids = ps_ids_by_location.get(&lb.id).cloned().unwrap_or_default();
                let rr_ids = rr_ids_by_location.get(&lb.id).cloned().unwrap_or_default();
                (lb, ar_ids, lr_ids, ps_ids, rr_ids)
            })
            .collect();

        let server_block_tuples: Vec<(
            crate::db::entities::server_block::Model,
            Vec<OverrideRef>,
            Vec<OverrideRef>,
            Vec<OverrideRef>,
            Vec<OverrideRef>,
        )> = server_blocks
            .into_iter()
            .map(|sb| {
                let ar_ids = server_ar_ids_by_server
                    .get(&sb.id)
                    .cloned()
                    .unwrap_or_default();
                let lb_ids = loc_block_ids_by_server
                    .get(&sb.id)
                    .cloned()
                    .unwrap_or_default();
                let ls_ids = log_setting_ids_by_server
                    .get(&sb.id)
                    .cloned()
                    .unwrap_or_default();
                let sc_ids: Vec<OverrideRef> = sb
                    .ssl_cert_id
                    .into_iter()
                    .map(|id| OverrideRef {
                        id,
                        override_of_id: None,
                    })
                    .collect();
                (sb, ar_ids, lb_ids, ls_ids, sc_ids)
            })
            .collect();

        let upstream_tuples: Vec<(crate::db::entities::upstream::Model, Vec<OverrideRef>)> =
            upstreams
                .into_iter()
                .map(|u| {
                    let lb_ids = loc_block_ids_by_upstream
                        .get(&u.id)
                        .cloned()
                        .unwrap_or_default();
                    (u, lb_ids)
                })
                .collect();

        // Combine server-level and location-level access rules into one flat list
        let all_access_rules: Vec<crate::db::entities::access_rule::Model> = {
            let mut ars =
                Vec::with_capacity(server_access_rules.len() + location_access_rules.len());
            ars.extend(server_access_rules);
            ars.extend(location_access_rules);
            ars
        };

        Ok(ProxyConfig {
            id: proxy_entity.id,
            name: proxy_entity.name,
            r#type: ProxyType::Nginx,
            description: proxy_entity.description,
            parent_config_id,
            server_blocks: server_block_tuples
                .into_iter()
                .map(|m| (m.0.id, m.into()))
                .collect(),
            upstreams: upstream_tuples
                .into_iter()
                .map(|m| (m.0.id, m.into()))
                .collect(),
            location_blocks: location_block_tuples
                .into_iter()
                .map(|m| (m.0.id, m.into()))
                .collect(),
            access_rules: all_access_rules
                .into_iter()
                .map(|m| (m.id, m.into()))
                .collect(),
            cache_zones: cache_zones.into_iter().map(|m| (m.id, m.into())).collect(),
            limit_rules: limit_rules.into_iter().map(|m| (m.id, m.into())).collect(),
            limit_zones: limit_zones.into_iter().map(|m| (m.id, m.into())).collect(),
            log_settings: log_settings.into_iter().map(|m| (m.id, m.into())).collect(),
            proxy_settings: proxy_settings
                .into_iter()
                .map(|m| (m.id, m.into()))
                .collect(),
            rewrite_rules: rewrite_rules
                .into_iter()
                .map(|m| (m.id, m.into()))
                .collect(),
            ssl_certificates: ssl_certificates
                .into_iter()
                .map(|m| (m.id, m.into()))
                .collect(),
        })
    }

    async fn get_merged_proxy_config(
        &self,
        proxy_id: uuid::Uuid,
    ) -> ProxyServiceResult<ProxyConfig> {
        let configs = self.get_proxy_raw_configs(proxy_id).await?;

        // configs is ordered [leaf, ..., root] (most specific first)
        // self.merge(other) means self overrides other
        // So start with leaf and merge each ancestor into it
        let mut iter = configs.into_iter();
        let mut merged = iter.next().ok_or(ProxyServiceError::ConfigNotFound)?;
        for config in iter {
            merged.merge(config);
        }
        Ok(merged)
    }
}