zebflow 0.10.10

//! Project DB connection management service.

use std::sync::Arc;

use rand::Rng;
use serde_json::json;

use crate::platform::adapters::data::DataAdapter;
use crate::platform::error::PlatformError;
use crate::platform::model::{
    ProjectDbConnection, ProjectDbConnectionListItem, ProjectDbConnectionTestResult,
    TestProjectDbConnectionRequest, UpsertProjectDbConnectionRequest, now_ts, slug_segment,
};
use crate::platform::sekejap;

/// Project-scoped DB connections stored in the metadata catalog.
pub struct DbConnectionService {
    data: Arc<dyn DataAdapter>,
}

impl DbConnectionService {
    /// Creates the DB connection service.
    pub fn new(data: Arc<dyn DataAdapter>) -> Self {
        Self { data }
    }

    /// Lists project DB connections (ensures the default connection exists).
    pub fn list_project_connections(
        &self,
        owner: &str,
        project: &str,
    ) -> Result<Vec<ProjectDbConnectionListItem>, PlatformError> {
        let owner = slug_segment(owner);
        let project = slug_segment(project);
        self.ensure_project_exists(&owner, &project)?;
        self.ensure_default_connection(&owner, &project)?;

        let mut items = self
            .data
            .list_project_db_connections(&owner, &project)?
            .into_iter()
            .map(|connection| ProjectDbConnectionListItem {
                connection_id: connection.connection_id,
                connection_slug: connection.connection_slug,
                connection_label: connection.connection_label,
                database_kind: connection.database_kind,
                credential_id: connection.credential_id,
                created_at: connection.created_at,
                updated_at: connection.updated_at,
            })
            .collect::<Vec<_>>();
        items.sort_by(|a, b| a.connection_slug.cmp(&b.connection_slug));
        Ok(items)
    }

    /// Resolves one project DB connection by slug.
    pub fn get_project_connection(
        &self,
        owner: &str,
        project: &str,
        connection_slug: &str,
    ) -> Result<Option<ProjectDbConnection>, PlatformError> {
        let owner = slug_segment(owner);
        let project = slug_segment(project);
        let connection_slug = slug_segment(connection_slug);
        self.ensure_project_exists(&owner, &project)?;
        self.ensure_default_connection(&owner, &project)?;
        if connection_slug.is_empty() {
            return Ok(None);
        }
        self.data
            .get_project_db_connection(&owner, &project, &connection_slug)
    }

    /// Creates or updates one project DB connection.
    pub fn upsert_project_connection(
        &self,
        owner: &str,
        project: &str,
        req: &UpsertProjectDbConnectionRequest,
    ) -> Result<ProjectDbConnection, PlatformError> {
        let owner = slug_segment(owner);
        let project = slug_segment(project);
        self.ensure_project_exists(&owner, &project)?;
        self.ensure_default_connection(&owner, &project)?;

        let connection_slug = slug_segment(&req.connection_slug);
        if connection_slug.is_empty() {
            return Err(PlatformError::new(
                "PLATFORM_DB_CONNECTION_INVALID",
                "connection slug must not be empty",
            ));
        }

        let database_kind = normalize_database_kind(&req.database_kind)?;
        let credential_id = normalize_optional_slug(req.credential_id.as_deref());
        self.validate_credential_binding(
            &owner,
            &project,
            &database_kind,
            credential_id.as_deref(),
        )?;

        let now = now_ts();
        let existing = self
            .data
            .get_project_db_connection(&owner, &project, &connection_slug)?;
        let created_at = existing.as_ref().map(|row| row.created_at).unwrap_or(now);
        let connection_id = existing
            .as_ref()
            .map(|row| row.connection_id.clone())
            .unwrap_or_else(generate_connection_id);
        let connection = ProjectDbConnection {
            owner: owner.clone(),
            project: project.clone(),
            connection_id,
            connection_slug: connection_slug.clone(),
            connection_label: if req.connection_label.trim().is_empty() {
                connection_slug.replace('-', " ")
            } else {
                req.connection_label.trim().to_string()
            },
            database_kind,
            credential_id,
            config: req.config.clone(),
            created_at,
            updated_at: now,
        };
        self.data.put_project_db_connection(&connection)?;
        Ok(connection)
    }

    /// Deletes one project DB connection by slug.
    pub fn delete_project_connection(
        &self,
        owner: &str,
        project: &str,
        connection_slug: &str,
    ) -> Result<(), PlatformError> {
        let owner = slug_segment(owner);
        let project = slug_segment(project);
        let connection_slug = slug_segment(connection_slug);
        self.ensure_project_exists(&owner, &project)?;
        if connection_slug.is_empty() {
            return Err(PlatformError::new(
                "PLATFORM_DB_CONNECTION_INVALID",
                "connection slug must not be empty",
            ));
        }
        if is_builtin_connection_slug(&connection_slug) {
            return Err(PlatformError::new(
                "PLATFORM_DB_CONNECTION_LOCKED",
                format!(
                    "built-in connection '{}' cannot be deleted",
                    connection_slug
                ),
            ));
        }
        self.data
            .delete_project_db_connection(&owner, &project, &connection_slug)
    }

    /// Tests one DB connection by actually connecting to the database.
    pub async fn test_project_connection(
        &self,
        owner: &str,
        project: &str,
        req: &TestProjectDbConnectionRequest,
    ) -> Result<ProjectDbConnectionTestResult, PlatformError> {
        let owner = slug_segment(owner);
        let project = slug_segment(project);
        self.ensure_project_exists(&owner, &project)?;
        self.ensure_default_connection(&owner, &project)?;

        let (database_kind, credential_id) = if let Some(slug) = req
            .connection_slug
            .as_deref()
            .map(slug_segment)
            .filter(|v| !v.is_empty())
        {
            let Some(connection) = self
                .data
                .get_project_db_connection(&owner, &project, &slug)?
            else {
                return Err(PlatformError::new(
                    "PLATFORM_DB_CONNECTION_MISSING",
                    format!("connection '{}' not found", slug),
                ));
            };
            (connection.database_kind, connection.credential_id)
        } else {
            let Some(kind) = req.database_kind.as_deref() else {
                return Err(PlatformError::new(
                    "PLATFORM_DB_CONNECTION_INVALID",
                    "database_kind is required when connection_slug is missing",
                ));
            };
            (
                normalize_database_kind(kind)?,
                normalize_optional_slug(req.credential_id.as_deref()),
            )
        };

        // Validate credential binding first
        self.validate_credential_binding(
            &owner,
            &project,
            &database_kind,
            credential_id.as_deref(),
        )?;

        // Now actually test the connection based on database kind
        match database_kind.as_str() {
            "postgresql" => {
                // Test PostgreSQL connection using sqlx (async)
                let credential = self
                    .data
                    .get_project_credential(&owner, &project, credential_id.as_ref().unwrap())?
                    .ok_or_else(|| {
                        PlatformError::new(
                            "PLATFORM_DB_CONNECTION_INVALID",
                            "credential not found".to_string(),
                        )
                    })?;

                let connect_options = build_postgres_connect_options(&credential.secret)?;

                // Connect using sqlx pool and run a smoke query.
                let pool = sqlx::postgres::PgPoolOptions::new()
                    .max_connections(1)
                    .acquire_timeout(std::time::Duration::from_secs(5))
                    .connect_with(connect_options)
                    .await
                    .map_err(|e| {
                        PlatformError::new(
                            "PLATFORM_DB_CONNECTION_FAILED",
                            format!("failed to connect to PostgreSQL: {}", e),
                        )
                    })?;

                // Execute a simple query to verify the connection
                let _result: (i32,) =
                    sqlx::query_as("SELECT 1")
                        .fetch_one(&pool)
                        .await
                        .map_err(|e| {
                            PlatformError::new(
                                "PLATFORM_DB_CONNECTION_QUERY_FAILED",
                                format!("connection established but query failed: {}", e),
                            )
                        })?;

                // Get the PostgreSQL version for additional info
                let version_row: (String,) = sqlx::query_as("SELECT version()")
                    .fetch_one(&pool)
                    .await
                    .map_err(|e| {
                        PlatformError::new(
                            "PLATFORM_DB_CONNECTION_QUERY_FAILED",
                            format!("failed to get version: {}", e),
                        )
                    })?;

                Ok(ProjectDbConnectionTestResult {
                    ok: true,
                    message: format!(
                        "Successfully connected to PostgreSQL. Server: {}",
                        version_row
                            .0
                            .split(' ')
                            .take(2)
                            .collect::<Vec<_>>()
                            .join(" ")
                    ),
                    details: json!({
                        "database_kind": database_kind,
                        "credential_id": credential_id,
                        "mode": "connection_test",
                        "server_version": version_row.0,
                        "test_query": "SELECT 1"
                    }),
                })
            }
            "mysql" => {
                // Test MySQL connection - requires mysql crate
                let credential = self
                    .data
                    .get_project_credential(&owner, &project, credential_id.as_ref().unwrap())?
                    .ok_or_else(|| {
                        PlatformError::new(
                            "PLATFORM_DB_CONNECTION_INVALID",
                            "credential not found".to_string(),
                        )
                    })?;

                let (host, port, _user, _password, database) =
                    extract_mysql_credential(&credential.secret)?;

                // Try to connect using simple_mysql approach
                // For now, return that mysql testing is not yet implemented
                // Once mysql crate is added, implement actual connection test
                Ok(ProjectDbConnectionTestResult {
                    ok: true,
                    message: format!(
                        "MySQL connection configuration is valid (host: {}:{}, database: {})",
                        host, port, database
                    ),
                    details: json!({
                        "database_kind": database_kind,
                        "credential_id": credential_id,
                        "mode": "connection_test",
                        "note": "MySQL connection test requires mysql crate to be added"
                    }),
                })
            }
            _ => {
                // For other database types, just validate config
                Ok(ProjectDbConnectionTestResult {
                    ok: true,
                    message: format!("{} connection configuration is valid", database_kind),
                    details: json!({
                        "database_kind": database_kind,
                        "credential_id": credential_id,
                        "mode": "validation"
                    }),
                })
            }
        }
    }

    fn ensure_default_connection(&self, owner: &str, project: &str) -> Result<(), PlatformError> {
        if self
            .data
            .get_project_db_connection(owner, project, "default")?
            .is_some()
            && self
                .data
                .get_project_db_connection(owner, project, sekejap::BUILTIN_CONNECTION_SLUG)?
                .is_some()
        {
            return Ok(());
        }
        let now = now_ts();
        if self
            .data
            .get_project_db_connection(owner, project, "default")?
            .is_none()
        {
            self.data.put_project_db_connection(&ProjectDbConnection {
                owner: owner.to_string(),
                project: project.to_string(),
                connection_id: generate_connection_id(),
                connection_slug: "default".to_string(),
                connection_label: "Default Data Store".to_string(),
                database_kind: "sqlite".to_string(),
                credential_id: None,
                config: json!({}),
                created_at: now,
                updated_at: now,
            })?;
        }
        if self
            .data
            .get_project_db_connection(owner, project, sekejap::BUILTIN_CONNECTION_SLUG)?
            .is_none()
        {
            self.data.put_project_db_connection(&ProjectDbConnection {
                owner: owner.to_string(),
                project: project.to_string(),
                connection_id: generate_connection_id(),
                connection_slug: sekejap::BUILTIN_CONNECTION_SLUG.to_string(),
                connection_label: sekejap::BUILTIN_CONNECTION_LABEL.to_string(),
                database_kind: sekejap::DB_KIND.to_string(),
                credential_id: None,
                config: json!({}),
                created_at: now,
                updated_at: now,
            })?;
        }
        Ok(())
    }

    fn validate_credential_binding(
        &self,
        owner: &str,
        project: &str,
        database_kind: &str,
        credential_id: Option<&str>,
    ) -> Result<(), PlatformError> {
        match database_kind {
            "sqlite" | "sekejap" => {
                if credential_id.is_some() {
                    return Err(PlatformError::new(
                        "PLATFORM_DB_CONNECTION_INVALID",
                        format!("{database_kind} connection must not bind credential_id"),
                    ));
                }
                Ok(())
            }
            "postgresql" => {
                let Some(credential_id) = credential_id else {
                    return Err(PlatformError::new(
                        "PLATFORM_DB_CONNECTION_INVALID",
                        "postgresql connection requires credential_id",
                    ));
                };
                let Some(credential) =
                    self.data
                        .get_project_credential(owner, project, credential_id)?
                else {
                    return Err(PlatformError::new(
                        "PLATFORM_DB_CONNECTION_INVALID",
                        format!("credential '{}' not found", credential_id),
                    ));
                };
                if credential.kind != "postgres" {
                    return Err(PlatformError::new(
                        "PLATFORM_DB_CONNECTION_INVALID",
                        format!(
                            "credential '{}' kind '{}' is not compatible with postgresql",
                            credential_id, credential.kind
                        ),
                    ));
                }
                Ok(())
            }
            "mysql" => {
                let Some(credential_id) = credential_id else {
                    return Err(PlatformError::new(
                        "PLATFORM_DB_CONNECTION_INVALID",
                        "mysql connection requires credential_id",
                    ));
                };
                let Some(credential) =
                    self.data
                        .get_project_credential(owner, project, credential_id)?
                else {
                    return Err(PlatformError::new(
                        "PLATFORM_DB_CONNECTION_INVALID",
                        format!("credential '{}' not found", credential_id),
                    ));
                };
                if credential.kind != "mysql" {
                    return Err(PlatformError::new(
                        "PLATFORM_DB_CONNECTION_INVALID",
                        format!(
                            "credential '{}' kind '{}' is not compatible with mysql",
                            credential_id, credential.kind
                        ),
                    ));
                }
                Ok(())
            }
            _ => Ok(()),
        }
    }

    fn ensure_project_exists(&self, owner: &str, project: &str) -> Result<(), PlatformError> {
        if self.data.get_project(owner, project)?.is_some() {
            return Ok(());
        }
        Err(PlatformError::new(
            "PLATFORM_PROJECT_MISSING",
            format!("project '{owner}/{project}' not found"),
        ))
    }
}

fn normalize_optional_slug(raw: Option<&str>) -> Option<String> {
    raw.map(slug_segment).filter(|value| !value.is_empty())
}

fn is_builtin_connection_slug(slug: &str) -> bool {
    slug == "default" || slug == sekejap::BUILTIN_CONNECTION_SLUG
}

fn generate_connection_id() -> String {
    let mut bytes = [0u8; 16];
    rand::rng().fill_bytes(&mut bytes);
    hex::encode(bytes)
}

fn normalize_database_kind(raw: &str) -> Result<String, PlatformError> {
    let normalized = slug_segment(raw);
    let kind = match normalized.as_str() {
        "postgres" | "postgresql" | "pg" => "postgresql",
        "mysql" => "mysql",
        "sqlite" => "sqlite",
        "sekejap" => "sekejap",
        "sqlserver" | "mssql" => "sqlserver",
        "mongodb" | "mongo" => "mongodb",
        "redis" => "redis",
        "qdrant" => "qdrant",
        "pinecone" => "pinecone",
        "chromadb" => "chromadb",
        "elasticsearch" | "elastic" => "elasticsearch",
        _ => {
            return Err(PlatformError::new(
                "PLATFORM_DB_CONNECTION_KIND_INVALID",
                format!("unsupported database kind '{}'", raw.trim()),
            ));
        }
    };
    Ok(kind.to_string())
}

/// Builds a PostgreSQL connection string from credential secret.
fn build_postgres_connect_options(
    secret: &serde_json::Value,
) -> Result<sqlx::postgres::PgConnectOptions, PlatformError> {
    let host = secret
        .get("host")
        .and_then(serde_json::Value::as_str)
        .ok_or_else(|| PlatformError::new("PLATFORM_DB_SECRET", "secret.host is required"))?;
    let port = secret
        .get("port")
        .and_then(|value| {
            value.as_u64().or_else(|| {
                value
                    .as_str()
                    .and_then(|raw| raw.trim().parse::<u64>().ok())
            })
        })
        .unwrap_or(5432);
    let port = u16::try_from(port).map_err(|_| {
        PlatformError::new("PLATFORM_DB_SECRET", "secret.port must be in 0..=65535")
    })?;
    let database = secret
        .get("database")
        .and_then(serde_json::Value::as_str)
        .ok_or_else(|| PlatformError::new("PLATFORM_DB_SECRET", "secret.database is required"))?;
    let user = secret
        .get("user")
        .and_then(serde_json::Value::as_str)
        .ok_or_else(|| PlatformError::new("PLATFORM_DB_SECRET", "secret.user is required"))?;
    let password = secret
        .get("password")
        .and_then(serde_json::Value::as_str)
        .unwrap_or("");
    Ok(sqlx::postgres::PgConnectOptions::new()
        .host(host)
        .port(port)
        .database(database)
        .username(user)
        .password(password))
}

/// Extracts MySQL connection parameters from credential secret.
fn extract_mysql_credential(
    secret: &serde_json::Value,
) -> Result<(String, u16, String, String, String), PlatformError> {
    let host = secret
        .get("host")
        .and_then(serde_json::Value::as_str)
        .unwrap_or("localhost");
    let port = secret
        .get("port")
        .and_then(|value| {
            value.as_u64().or_else(|| {
                value
                    .as_str()
                    .and_then(|raw| raw.trim().parse::<u64>().ok())
            })
        })
        .unwrap_or(3306) as u16;
    let user = secret
        .get("user")
        .and_then(serde_json::Value::as_str)
        .ok_or_else(|| PlatformError::new("PLATFORM_DB_SECRET", "secret.user is required"))?;
    let password = secret
        .get("password")
        .and_then(serde_json::Value::as_str)
        .unwrap_or("");
    let database = secret
        .get("database")
        .and_then(serde_json::Value::as_str)
        .unwrap_or("");
    Ok((
        host.to_string(),
        port,
        user.to_string(),
        password.to_string(),
        database.to_string(),
    ))
}