photoncloud-monorepo/prismnet/crates/prismnet-server/src/ovn/client.rs

use std::sync::Arc;

use prismnet_types::{DhcpOptions, Port, PortId, SecurityGroupId, SecurityGroupRule, SecurityGroupRuleId, VpcId};
use tokio::process::Command;
use tokio::sync::Mutex;

use crate::ovn::mock::MockOvnState;

/// OVN client mode
#[derive(Debug, Clone)]
pub enum OvnMode {
    Real { nb_addr: String },
    Mock(Arc<Mutex<MockOvnState>>),
}

/// OVN client error
#[derive(Debug, thiserror::Error)]
pub enum OvnError {
    #[error("OVN command failed: {0}")]
    Command(String),
    #[error("Invalid argument: {0}")]
    InvalidArgument(String),
}

pub type OvnResult<T> = std::result::Result<T, OvnError>;

/// Lightweight OVN client with mock/real modes
#[derive(Clone)]
pub struct OvnClient {
    mode: OvnMode,
}

impl OvnClient {
    /// Build an OVN client from environment variables (default: mock)
    /// - NOVANET_OVN_MODE: "mock" (default) or "real"
    /// - NOVANET_OVN_NB_ADDR: ovsdb northbound address (real mode only)
    pub fn from_env() -> OvnResult<Self> {
        let mode = std::env::var("NOVANET_OVN_MODE").unwrap_or_else(|_| "mock".to_string());
        match mode.to_lowercase().as_str() {
            "mock" => Ok(Self::new_mock()),
            "real" => {
                let nb_addr = std::env::var("NOVANET_OVN_NB_ADDR")
                    .unwrap_or_else(|_| "tcp:127.0.0.1:6641".to_string());
                Ok(Self::new_real(nb_addr))
            }
            other => Err(OvnError::InvalidArgument(format!(
                "Unknown OVN mode: {}",
                other
            ))),
        }
    }

    pub fn new_mock() -> Self {
        Self {
            mode: OvnMode::Mock(Arc::new(Mutex::new(MockOvnState::new()))),
        }
    }

    pub fn new_real(nb_addr: impl Into<String>) -> Self {
        Self {
            mode: OvnMode::Real {
                nb_addr: nb_addr.into(),
            },
        }
    }

    /// Expose mock state for tests
    pub fn mock_state(&self) -> Option<Arc<Mutex<MockOvnState>>> {
        match &self.mode {
            OvnMode::Mock(state) => Some(state.clone()),
            _ => None,
        }
    }

    fn logical_switch_name(vpc_id: &VpcId) -> String {
        format!("vpc-{}", vpc_id)
    }

    fn logical_port_name(port_id: &PortId) -> String {
        format!("port-{}", port_id)
    }

    fn acl_key(rule_id: &SecurityGroupRuleId) -> String {
        format!("acl-{}", rule_id)
    }

    async fn run_nbctl(&self, args: Vec<String>) -> OvnResult<()> {
        let nb_addr = match &self.mode {
            OvnMode::Real { nb_addr } => nb_addr,
            _ => {
                return Err(OvnError::InvalidArgument(
                    "nbctl invocation only valid in real mode".to_string(),
                ))
            }
        };

        let output = Command::new("ovn-nbctl")
            .args(["--db", nb_addr])
            .args(args.iter().map(String::as_str))
            .output()
            .await
            .map_err(|e| OvnError::Command(format!("failed to run ovn-nbctl: {}", e)))?;

        if output.status.success() {
            Ok(())
        } else {
            let stderr = String::from_utf8_lossy(&output.stderr).to_string();
            Err(OvnError::Command(format!(
                "ovn-nbctl exit code {}: {}",
                output.status, stderr
            )))
        }
    }

    pub async fn create_logical_switch(&self, vpc_id: &VpcId, cidr: &str) -> OvnResult<()> {
        match &self.mode {
            OvnMode::Mock(state) => {
                let mut guard = state.lock().await;
                guard.create_logical_switch(*vpc_id, cidr.to_string());
                Ok(())
            }
            OvnMode::Real { .. } => {
                let name = Self::logical_switch_name(vpc_id);
                self.run_nbctl(vec!["ls-add".into(), name.clone()]).await?;
                // Store CIDR for reference (best-effort; ignore errors)
                let _ = self
                    .run_nbctl(vec![
                        "set".into(),
                        "Logical_Switch".into(),
                        name,
                        format!("other_config:subnet={}", cidr),
                    ])
                    .await;
                Ok(())
            }
        }
    }

    pub async fn delete_logical_switch(&self, vpc_id: &VpcId) -> OvnResult<()> {
        match &self.mode {
            OvnMode::Mock(state) => {
                let mut guard = state.lock().await;
                guard.delete_logical_switch(vpc_id);
                Ok(())
            }
            OvnMode::Real { .. } => {
                let name = Self::logical_switch_name(vpc_id);
                self.run_nbctl(vec!["ls-del".into(), name]).await
            }
        }
    }

    pub async fn create_logical_switch_port(
        &self,
        port: &Port,
        vpc_id: &VpcId,
        ip_address: &str,
    ) -> OvnResult<()> {
        match &self.mode {
            OvnMode::Mock(state) => {
                let mut guard = state.lock().await;
                guard.create_logical_switch_port(
                    port.id,
                    *vpc_id,
                    port.mac_address.clone(),
                    ip_address.to_string(),
                );
                Ok(())
            }
            OvnMode::Real { .. } => {
                let ls_name = Self::logical_switch_name(vpc_id);
                let lsp_name = Self::logical_port_name(&port.id);
                self.run_nbctl(vec!["lsp-add".into(), ls_name.clone(), lsp_name.clone()])
                    .await?;

                let address = format!("{} {}", port.mac_address, ip_address);
                self.run_nbctl(vec![
                    "set".into(),
                    "Logical_Switch_Port".into(),
                    lsp_name,
                    format!("addresses=\"{}\"", address),
                ])
                .await?;
                Ok(())
            }
        }
    }

    pub async fn delete_logical_switch_port(&self, port_id: &PortId) -> OvnResult<()> {
        match &self.mode {
            OvnMode::Mock(state) => {
                let mut guard = state.lock().await;
                guard.delete_logical_switch_port(port_id);
                Ok(())
            }
            OvnMode::Real { .. } => {
                let lsp_name = Self::logical_port_name(port_id);
                self.run_nbctl(vec!["lsp-del".into(), lsp_name]).await
            }
        }
    }

    pub async fn create_acl(
        &self,
        sg_id: &SecurityGroupId,
        rule: &SecurityGroupRule,
        logical_switch: &VpcId,
        match_expr: &str,
        priority: u16,
    ) -> OvnResult<String> {
        let key = Self::acl_key(&rule.id);
        match &self.mode {
            OvnMode::Mock(state) => {
                let mut guard = state.lock().await;
                guard.create_acl(
                    key.clone(),
                    *sg_id,
                    rule.clone(),
                    *logical_switch,
                    match_expr.to_string(),
                );
                Ok(key)
            }
            OvnMode::Real { .. } => {
                let ls_name = Self::logical_switch_name(logical_switch);
                let direction = direction_to_ovn(rule);

                // ovn-nbctl acl-add <switch> <direction> <priority> <match> <action>
                self.run_nbctl(vec![
                    "acl-add".into(),
                    ls_name,
                    direction,
                    priority.to_string(),
                    match_expr.to_string(),
                    "allow-related".into(),
                ])
                .await?;
                Ok(key)
            }
        }
    }

    pub async fn delete_acl(&self, rule_id: &SecurityGroupRuleId) -> OvnResult<()> {
        let key = Self::acl_key(rule_id);
        match &self.mode {
            OvnMode::Mock(state) => {
                let mut guard = state.lock().await;
                guard.delete_acl(&key);
                Ok(())
            }
            OvnMode::Real { .. } => {
                // Best-effort deletion by external-id match (placeholder)
                let _ = self
                    .run_nbctl(vec![
                        "--".into(),
                        "find".into(),
                        "ACL".into(),
                        format!("name={}", key),
                        "--delete".into(),
                        "ACL".into(),
                        "uuid".into(),
                    ])
                    .await;
                Ok(())
            }
        }
    }

    /// Create DHCP options in OVN for a subnet
    pub async fn create_dhcp_options(
        &self,
        cidr: &str,
        options: &DhcpOptions,
    ) -> OvnResult<String> {
        match &self.mode {
            OvnMode::Mock(state) => {
                let mut guard = state.lock().await;
                let uuid = guard.create_dhcp_options(cidr.to_string(), options.clone());
                Ok(uuid)
            }
            OvnMode::Real { nb_addr } => {
                // Command: ovn-nbctl dhcp-options-create <cidr>
                let output = Command::new("ovn-nbctl")
                    .args(["--db", nb_addr])
                    .args(["dhcp-options-create", cidr])
                    .output()
                    .await
                    .map_err(|e| OvnError::Command(format!("failed to run ovn-nbctl: {}", e)))?;

                if !output.status.success() {
                    let stderr = String::from_utf8_lossy(&output.stderr).to_string();
                    return Err(OvnError::Command(format!(
                        "dhcp-options-create failed: {}",
                        stderr
                    )));
                }

                let uuid = String::from_utf8_lossy(&output.stdout).trim().to_string();

                // Set DHCP options
                let mut opts = Vec::new();

                if let Some(router) = &options.router {
                    opts.push(format!("router={}", router));
                }

                if !options.dns_servers.is_empty() {
                    opts.push(format!(
                        "dns_server={{{}}}",
                        options.dns_servers.join(",")
                    ));
                }

                opts.push(format!("lease_time={}", options.lease_time));

                if let Some(domain) = &options.domain_name {
                    opts.push(format!("domain_name={}", domain));
                }

                // Command: ovn-nbctl dhcp-options-set-options <uuid> <options>
                let opts_str = opts.join(" ");
                self.run_nbctl(vec![
                    "dhcp-options-set-options".into(),
                    uuid.clone(),
                    opts_str,
                ])
                .await?;

                Ok(uuid)
            }
        }
    }

    /// Delete DHCP options from OVN
    pub async fn delete_dhcp_options(&self, uuid: &str) -> OvnResult<()> {
        match &self.mode {
            OvnMode::Mock(state) => {
                let mut guard = state.lock().await;
                guard.delete_dhcp_options(uuid);
                Ok(())
            }
            OvnMode::Real { .. } => {
                self.run_nbctl(vec!["dhcp-options-del".into(), uuid.to_string()])
                    .await
            }
        }
    }

    /// Associate DHCP options with a logical switch port
    pub async fn set_lsp_dhcp_options(&self, lsp_name: &str, dhcp_uuid: &str) -> OvnResult<()> {
        match &self.mode {
            OvnMode::Mock(state) => {
                let mut guard = state.lock().await;
                guard.set_lsp_dhcp_options(lsp_name.to_string(), dhcp_uuid.to_string());
                Ok(())
            }
            OvnMode::Real { .. } => {
                self.run_nbctl(vec![
                    "lsp-set-dhcpv4-options".into(),
                    lsp_name.to_string(),
                    dhcp_uuid.to_string(),
                ])
                .await
            }
        }
    }

    /// Create a logical router
    pub async fn create_logical_router(&self, name: &str) -> OvnResult<String> {
        match &self.mode {
            OvnMode::Mock(state) => {
                let mut guard = state.lock().await;
                let id = guard.create_router(name.to_string());
                Ok(id)
            }
            OvnMode::Real { nb_addr } => {
                // Command: ovn-nbctl lr-add <router-name>
                self.run_nbctl(vec!["lr-add".into(), name.to_string()])
                    .await?;

                // Get the UUID of the created router
                let output = Command::new("ovn-nbctl")
                    .args(["--db", nb_addr])
                    .args(["--columns=_uuid", "--bare", "find", "Logical_Router"])
                    .arg(format!("name={}", name))
                    .output()
                    .await
                    .map_err(|e| OvnError::Command(format!("failed to run ovn-nbctl: {}", e)))?;

                if !output.status.success() {
                    let stderr = String::from_utf8_lossy(&output.stderr).to_string();
                    return Err(OvnError::Command(format!("lr-add query failed: {}", stderr)));
                }

                let uuid = String::from_utf8_lossy(&output.stdout).trim().to_string();
                Ok(uuid)
            }
        }
    }

    /// Delete a logical router
    pub async fn delete_logical_router(&self, router_id: &str) -> OvnResult<()> {
        match &self.mode {
            OvnMode::Mock(state) => {
                let mut guard = state.lock().await;
                guard.delete_router(router_id);
                Ok(())
            }
            OvnMode::Real { .. } => {
                self.run_nbctl(vec!["lr-del".into(), router_id.to_string()])
                    .await
            }
        }
    }

    /// Add a router port connecting a router to a logical switch
    /// Returns the router port ID
    pub async fn add_router_port(
        &self,
        router_id: &str,
        switch_id: &VpcId,
        cidr: &str,
        mac: &str,
    ) -> OvnResult<String> {
        match &self.mode {
            OvnMode::Mock(state) => {
                let mut guard = state.lock().await;
                let port_id = guard.add_router_port(
                    router_id.to_string(),
                    *switch_id,
                    cidr.to_string(),
                    mac.to_string(),
                );
                Ok(port_id)
            }
            OvnMode::Real { .. } => {
                // Generate unique port names
                let router_port_name = format!("rtr-port-{}", uuid::Uuid::new_v4());
                let switch_port_name = format!("lsp-rtr-{}", uuid::Uuid::new_v4());

                // Extract the IP address from CIDR for the router port
                let ip_with_prefix = cidr;

                // Create logical router port
                // ovn-nbctl lrp-add <router> <port-name> <mac> <network>
                self.run_nbctl(vec![
                    "lrp-add".into(),
                    router_id.to_string(),
                    router_port_name.clone(),
                    mac.to_string(),
                    ip_with_prefix.to_string(),
                ])
                .await?;

                // Create the corresponding switch port
                let ls_name = Self::logical_switch_name(switch_id);

                // ovn-nbctl lsp-add <switch> <port-name>
                self.run_nbctl(vec![
                    "lsp-add".into(),
                    ls_name,
                    switch_port_name.clone(),
                ])
                .await?;

                // ovn-nbctl lsp-set-type <port> router
                self.run_nbctl(vec![
                    "lsp-set-type".into(),
                    switch_port_name.clone(),
                    "router".into(),
                ])
                .await?;

                // ovn-nbctl lsp-set-addresses <port> router
                self.run_nbctl(vec![
                    "lsp-set-addresses".into(),
                    switch_port_name.clone(),
                    "router".into(),
                ])
                .await?;

                // ovn-nbctl lsp-set-options <port> router-port=<router-port>
                self.run_nbctl(vec![
                    "lsp-set-options".into(),
                    switch_port_name,
                    format!("router-port={}", router_port_name),
                ])
                .await?;

                Ok(router_port_name)
            }
        }
    }

    /// Configure SNAT on a logical router
    pub async fn configure_snat(
        &self,
        router_id: &str,
        external_ip: &str,
        logical_ip_cidr: &str,
    ) -> OvnResult<()> {
        match &self.mode {
            OvnMode::Mock(state) => {
                let mut guard = state.lock().await;
                guard.configure_snat(
                    router_id.to_string(),
                    external_ip.to_string(),
                    logical_ip_cidr.to_string(),
                );
                Ok(())
            }
            OvnMode::Real { .. } => {
                // ovn-nbctl lr-nat-add <router> snat <external-ip> <logical-ip-cidr>
                self.run_nbctl(vec![
                    "lr-nat-add".into(),
                    router_id.to_string(),
                    "snat".into(),
                    external_ip.to_string(),
                    logical_ip_cidr.to_string(),
                ])
                .await
            }
        }
    }
}

fn direction_to_ovn(rule: &SecurityGroupRule) -> String {
    match rule.direction {
        prismnet_types::RuleDirection::Ingress => "to-lport".to_string(),
        prismnet_types::RuleDirection::Egress => "from-lport".to_string(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use prismnet_types::{RuleDirection, SecurityGroupRule, Vpc};

    #[tokio::test]
    async fn mock_logical_switch_and_port_lifecycle() {
        let client = OvnClient::new_mock();
        let vpc = Vpc::new("test", "org", "proj", "10.0.0.0/16");
        client
            .create_logical_switch(&vpc.id, &vpc.cidr_block)
            .await
            .unwrap();

        let mut port = Port::new("p1", prismnet_types::SubnetId::new());
        port.ip_address = Some("10.0.0.5".to_string());

        client
            .create_logical_switch_port(&port, &vpc.id, port.ip_address.as_ref().unwrap())
            .await
            .unwrap();

        let state = client.mock_state().unwrap();
        let guard = state.lock().await;
        assert!(guard.has_logical_switch(&vpc.id));
        assert!(guard.port_attached(&port.id));
    }

    #[tokio::test]
    async fn mock_acl_tracks_rule() {
        let client = OvnClient::new_mock();
        let vpc = Vpc::new("test", "org", "proj", "10.0.0.0/16");
        client
            .create_logical_switch(&vpc.id, &vpc.cidr_block)
            .await
            .unwrap();

        let mut rule = SecurityGroupRule::new(
            SecurityGroupId::new(),
            RuleDirection::Ingress,
            Default::default(),
        );
        rule.remote_cidr = Some("0.0.0.0/0".to_string());

        let key = client
            .create_acl(
                &rule.security_group_id,
                &rule,
                &vpc.id,
                "ip4 && ip4.src == 0.0.0.0/0",
                1000,
            )
            .await
            .unwrap();
        let state = client.mock_state().unwrap();
        let guard = state.lock().await;
        assert!(guard.acl_exists(&key));
    }

    #[tokio::test]
    async fn mock_deletions_remove_state() {
        let client = OvnClient::new_mock();
        let vpc = Vpc::new("test", "org", "proj", "10.0.0.0/16");
        client
            .create_logical_switch(&vpc.id, &vpc.cidr_block)
            .await
            .unwrap();
        client.delete_logical_switch(&vpc.id).await.unwrap();
        let state = client.mock_state().unwrap();
        let guard = state.lock().await;
        assert!(!guard.has_logical_switch(&vpc.id));
    }

    #[tokio::test]
    async fn test_dhcp_options_lifecycle() {
        let client = OvnClient::new_mock();

        let opts = DhcpOptions {
            cidr: "192.168.1.0/24".to_string(),
            router: Some("192.168.1.1".to_string()),
            dns_servers: vec!["8.8.8.8".to_string(), "8.8.4.4".to_string()],
            lease_time: 3600,
            domain_name: Some("example.com".to_string()),
        };

        // Create DHCP options
        let uuid = client
            .create_dhcp_options("192.168.1.0/24", &opts)
            .await
            .unwrap();
        assert!(!uuid.is_empty());
        assert!(uuid.starts_with("dhcp-"));

        // Verify it was created
        let state = client.mock_state().unwrap();
        let guard = state.lock().await;
        assert!(guard.dhcp_options_exists(&uuid));
        drop(guard);

        // Delete DHCP options
        client.delete_dhcp_options(&uuid).await.unwrap();

        // Verify it was deleted
        let guard = state.lock().await;
        assert!(!guard.dhcp_options_exists(&uuid));
    }

    #[tokio::test]
    async fn test_dhcp_options_default() {
        let opts = DhcpOptions::default();
        assert_eq!(opts.lease_time, 86400);
        assert_eq!(opts.dns_servers, vec!["8.8.8.8"]);
    }

    #[tokio::test]
    async fn test_lsp_dhcp_binding() {
        let client = OvnClient::new_mock();
        let vpc = Vpc::new("test", "org", "proj", "10.0.0.0/16");

        // Create logical switch
        client
            .create_logical_switch(&vpc.id, &vpc.cidr_block)
            .await
            .unwrap();

        // Create DHCP options
        let opts = DhcpOptions {
            cidr: vpc.cidr_block.clone(),
            router: Some("10.0.0.1".to_string()),
            dns_servers: vec!["8.8.8.8".to_string()],
            lease_time: 86400,
            domain_name: None,
        };

        let dhcp_uuid = client
            .create_dhcp_options(&vpc.cidr_block, &opts)
            .await
            .unwrap();

        // Create a port
        let mut port = Port::new("p1", prismnet_types::SubnetId::new());
        port.ip_address = Some("10.0.0.5".to_string());

        client
            .create_logical_switch_port(&port, &vpc.id, port.ip_address.as_ref().unwrap())
            .await
            .unwrap();

        // Bind DHCP options to port
        let lsp_name = format!("port-{}", port.id);
        client
            .set_lsp_dhcp_options(&lsp_name, &dhcp_uuid)
            .await
            .unwrap();

        // Verify binding
        let state = client.mock_state().unwrap();
        let guard = state.lock().await;
        assert!(guard.port_has_dhcp(&lsp_name));
    }

    // Router tests
    #[tokio::test]
    async fn test_router_create_and_delete() {
        let client = OvnClient::new_mock();

        // Create a router
        let router_id = client
            .create_logical_router("test-router")
            .await
            .unwrap();

        assert!(!router_id.is_empty());
        assert!(router_id.starts_with("router-"));

        // Verify it exists
        let state = client.mock_state().unwrap();
        let guard = state.lock().await;
        assert!(guard.router_exists(&router_id));
        drop(guard);

        // Delete the router
        client.delete_logical_router(&router_id).await.unwrap();

        // Verify it's gone
        let guard = state.lock().await;
        assert!(!guard.router_exists(&router_id));
    }

    #[tokio::test]
    async fn test_router_port_attachment() {
        let client = OvnClient::new_mock();
        let vpc = Vpc::new("test", "org", "proj", "10.0.0.0/16");

        // Create logical switch
        client
            .create_logical_switch(&vpc.id, &vpc.cidr_block)
            .await
            .unwrap();

        // Create router
        let router_id = client
            .create_logical_router("test-router")
            .await
            .unwrap();

        // Add router port
        let mac = "02:00:00:00:00:01";
        let cidr = "10.0.0.1/24";
        let port_id = client
            .add_router_port(&router_id, &vpc.id, cidr, mac)
            .await
            .unwrap();

        assert!(!port_id.is_empty());
        assert!(port_id.starts_with("rtr-port-"));

        // Verify port exists
        let state = client.mock_state().unwrap();
        let guard = state.lock().await;
        assert!(guard.router_port_exists(&port_id));
        assert_eq!(guard.get_router_port_count(&router_id), 1);
    }

    #[tokio::test]
    async fn test_snat_configuration() {
        let client = OvnClient::new_mock();

        // Create router
        let router_id = client
            .create_logical_router("test-router")
            .await
            .unwrap();

        // Configure SNAT
        let external_ip = "203.0.113.10";
        let logical_ip_cidr = "10.0.0.0/24";
        client
            .configure_snat(&router_id, external_ip, logical_ip_cidr)
            .await
            .unwrap();

        // Verify SNAT rule exists
        let state = client.mock_state().unwrap();
        let guard = state.lock().await;
        assert!(guard.snat_rule_exists(&router_id, external_ip));
    }

    #[tokio::test]
    async fn test_router_deletion_cascades() {
        let client = OvnClient::new_mock();
        let vpc = Vpc::new("test", "org", "proj", "10.0.0.0/16");

        // Create logical switch
        client
            .create_logical_switch(&vpc.id, &vpc.cidr_block)
            .await
            .unwrap();

        // Create router
        let router_id = client
            .create_logical_router("test-router")
            .await
            .unwrap();

        // Add router port
        let mac = "02:00:00:00:00:01";
        let cidr = "10.0.0.1/24";
        let port_id = client
            .add_router_port(&router_id, &vpc.id, cidr, mac)
            .await
            .unwrap();

        // Configure SNAT
        let external_ip = "203.0.113.10";
        let logical_ip_cidr = "10.0.0.0/24";
        client
            .configure_snat(&router_id, external_ip, logical_ip_cidr)
            .await
            .unwrap();

        // Delete router
        client.delete_logical_router(&router_id).await.unwrap();

        // Verify everything is cleaned up
        let state = client.mock_state().unwrap();
        let guard = state.lock().await;
        assert!(!guard.router_exists(&router_id));
        assert!(!guard.router_port_exists(&port_id));
        assert!(!guard.snat_rule_exists(&router_id, external_ip));
    }

    #[tokio::test]
    async fn test_multiple_router_ports() {
        let client = OvnClient::new_mock();
        let vpc1 = Vpc::new("test1", "org", "proj", "10.0.0.0/16");
        let vpc2 = Vpc::new("test2", "org", "proj", "10.1.0.0/16");

        // Create logical switches
        client
            .create_logical_switch(&vpc1.id, &vpc1.cidr_block)
            .await
            .unwrap();
        client
            .create_logical_switch(&vpc2.id, &vpc2.cidr_block)
            .await
            .unwrap();

        // Create router
        let router_id = client
            .create_logical_router("test-router")
            .await
            .unwrap();

        // Add router ports to both switches
        let port1_id = client
            .add_router_port(&router_id, &vpc1.id, "10.0.0.1/24", "02:00:00:00:00:01")
            .await
            .unwrap();

        let port2_id = client
            .add_router_port(&router_id, &vpc2.id, "10.1.0.1/24", "02:00:00:00:00:02")
            .await
            .unwrap();

        // Verify both ports exist
        let state = client.mock_state().unwrap();
        let guard = state.lock().await;
        assert!(guard.router_port_exists(&port1_id));
        assert!(guard.router_port_exists(&port2_id));
        assert_eq!(guard.get_router_port_count(&router_id), 2);
    }

    #[tokio::test]
    async fn test_full_vpc_router_snat_workflow() {
        let client = OvnClient::new_mock();
        let vpc = Vpc::new("test", "org", "proj", "10.0.0.0/16");

        // Step 1: Create VPC (logical switch)
        client
            .create_logical_switch(&vpc.id, &vpc.cidr_block)
            .await
            .unwrap();

        // Step 2: Create router
        let router_id = client
            .create_logical_router("vpc-router")
            .await
            .unwrap();

        // Step 3: Attach router to switch
        let mac = "02:00:00:00:00:01";
        let gateway_cidr = "10.0.0.1/24";
        let port_id = client
            .add_router_port(&router_id, &vpc.id, gateway_cidr, mac)
            .await
            .unwrap();

        // Step 4: Configure SNAT for outbound traffic
        let external_ip = "203.0.113.10";
        let internal_cidr = "10.0.0.0/24";
        client
            .configure_snat(&router_id, external_ip, internal_cidr)
            .await
            .unwrap();

        // Verify the complete setup
        let state = client.mock_state().unwrap();
        let guard = state.lock().await;

        // Check switch exists
        assert!(guard.has_logical_switch(&vpc.id));

        // Check router exists
        assert!(guard.router_exists(&router_id));

        // Check router port exists
        assert!(guard.router_port_exists(&port_id));
        assert_eq!(guard.get_router_port_count(&router_id), 1);

        // Check SNAT rule exists
        assert!(guard.snat_rule_exists(&router_id, external_ip));
    }

    #[tokio::test]
    async fn test_multiple_snat_rules() {
        let client = OvnClient::new_mock();

        // Create router
        let router_id = client
            .create_logical_router("test-router")
            .await
            .unwrap();

        // Configure multiple SNAT rules
        client
            .configure_snat(&router_id, "203.0.113.10", "10.0.0.0/24")
            .await
            .unwrap();

        client
            .configure_snat(&router_id, "203.0.113.11", "10.1.0.0/24")
            .await
            .unwrap();

        // Verify both SNAT rules exist
        let state = client.mock_state().unwrap();
        let guard = state.lock().await;
        assert!(guard.snat_rule_exists(&router_id, "203.0.113.10"));
        assert!(guard.snat_rule_exists(&router_id, "203.0.113.11"));
    }
}