uri-register 0.3.0

// Copyright TELICENT LTD
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Python bindings for uri-register using PyO3

use crate::{CacheStrategy, PostgresUriRegister, UriService};
use pyo3::exceptions::{PyRuntimeError, PyValueError};
use pyo3::prelude::*;
use std::collections::HashMap;
use tokio::runtime::Runtime;

/// Python wrapper for PostgresUriRegister
#[pyclass(name = "UriRegister")]
struct PyUriRegister {
    inner: PostgresUriRegister,
    rt: Runtime,
}

/// Helper function to parse cache strategy string
fn parse_cache_strategy(cache_strategy: Option<String>) -> PyResult<Option<CacheStrategy>> {
    if let Some(strategy_str) = cache_strategy {
        let strategy = match strategy_str.to_lowercase().as_str() {
            "moka" => CacheStrategy::Moka,
            "lru" => CacheStrategy::Lru,
            _ => {
                return Err(PyValueError::new_err(format!(
                    "Invalid cache_strategy '{}'. Must be 'moka' or 'lru'",
                    strategy_str
                )))
            }
        };
        Ok(Some(strategy))
    } else {
        Ok(None) // Default to Moka
    }
}

/// Helper function to validate connection parameters
fn validate_params(database_url: &str, max_connections: u32, cache_size: usize) -> PyResult<()> {
    if database_url.is_empty() {
        return Err(PyValueError::new_err("database_url cannot be empty"));
    }
    if max_connections == 0 {
        return Err(PyValueError::new_err(
            "max_connections must be greater than 0",
        ));
    }
    if max_connections > 10_000 {
        return Err(PyValueError::new_err(
            "max_connections must be 10000 or less",
        ));
    }
    if cache_size == 0 {
        return Err(PyValueError::new_err("cache_size must be greater than 0"));
    }
    Ok(())
}

#[pymethods]
impl PyUriRegister {
    /// Create a new URI register connected to PostgreSQL (synchronous)
    ///
    /// Args:
    ///     database_url: PostgreSQL connection string (e.g., "postgres://user:password@host:port/database")
    ///     table_name: Name of the database table to use (default: "uri_register")
    ///     max_connections: Maximum number of connections in the pool (recommended: 10-50)
    ///     cache_size: Number of URI-to-ID mappings to cache in memory (recommended: 1,000-100,000)
    ///     cache_strategy: Cache strategy - "moka" (W-TinyLFU, default and recommended) or "lru"
    ///     use_tls: Whether to use TLS for database connections (default: False)
    ///     ca_cert_path: Path to a PEM-encoded CA certificate file for private/internal CAs (default: None)
    ///
    /// Returns:
    ///     UriRegister: A new URI register instance
    ///
    /// Example:
    ///     >>> register = UriRegister("postgres://localhost/mydb", "uri_register", 20, 10000)
    ///     >>> # With private CA:
    ///     >>> register = UriRegister("postgres://db.internal/mydb", "uri_register", 20, 10000,
    ///     ...     use_tls=True, ca_cert_path="/etc/ssl/certs/internal-ca.pem")
    #[new]
    #[pyo3(signature = (database_url, table_name, max_connections, cache_size, cache_strategy=None, use_tls=None, ca_cert_path=None))]
    fn new(
        database_url: String,
        table_name: String,
        max_connections: u32,
        cache_size: usize,
        cache_strategy: Option<String>,
        use_tls: Option<bool>,
        ca_cert_path: Option<String>,
    ) -> PyResult<Self> {
        validate_params(&database_url, max_connections, cache_size)?;
        let cache_strat = parse_cache_strategy(cache_strategy)?;

        let rt = Runtime::new()
            .map_err(|e| PyRuntimeError::new_err(format!("Failed to create runtime: {}", e)))?;

        let inner = rt
            .block_on(PostgresUriRegister::new_with_cache_strategy(
                &database_url,
                &table_name,
                max_connections,
                cache_size,
                cache_strat,
                use_tls,
                ca_cert_path.as_deref(),
            ))
            .map_err(|e| PyRuntimeError::new_err(format!("Failed to connect: {}", e)))?;

        Ok(Self { inner, rt })
    }

    /// Create a new URI register connected to PostgreSQL (asynchronous)
    ///
    /// Args:
    ///     database_url: PostgreSQL connection string (e.g., "postgres://user:password@host:port/database")
    ///     table_name: Name of the database table to use (default: "uri_register")
    ///     max_connections: Maximum number of connections in the pool (recommended: 10-50)
    ///     cache_size: Number of URI-to-ID mappings to cache in memory (recommended: 1,000-100,000)
    ///     cache_strategy: Cache strategy - "moka" (W-TinyLFU, default and recommended) or "lru"
    ///     use_tls: Whether to use TLS for database connections (default: False)
    ///     ca_cert_path: Path to a PEM-encoded CA certificate file for private/internal CAs (default: None)
    ///
    /// Returns:
    ///     UriRegister: A new URI register instance
    ///
    /// Example:
    ///     >>> register = await UriRegister.new_async("postgres://localhost/mydb", "uri_register", 20, 10000)
    #[staticmethod]
    #[pyo3(signature = (database_url, table_name, max_connections, cache_size, cache_strategy=None, use_tls=None, ca_cert_path=None))]
    #[allow(clippy::too_many_arguments)]
    fn new_async<'py>(
        py: Python<'py>,
        database_url: String,
        table_name: String,
        max_connections: u32,
        cache_size: usize,
        cache_strategy: Option<String>,
        use_tls: Option<bool>,
        ca_cert_path: Option<String>,
    ) -> PyResult<Bound<'py, PyAny>> {
        validate_params(&database_url, max_connections, cache_size)?;
        let cache_strat = parse_cache_strategy(cache_strategy)?;

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let rt = Runtime::new()
                .map_err(|e| PyRuntimeError::new_err(format!("Failed to create runtime: {}", e)))?;

            let inner = PostgresUriRegister::new_with_cache_strategy(
                &database_url,
                &table_name,
                max_connections,
                cache_size,
                cache_strat,
                use_tls,
                ca_cert_path.as_deref(),
            )
            .await
            .map_err(|e| PyRuntimeError::new_err(format!("Failed to connect: {}", e)))?;

            Ok(PyUriRegister { inner, rt })
        })
    }

    // ==================== Synchronous methods ====================

    /// Register a single URI and return its ID (synchronous)
    ///
    /// Args:
    ///     uri: The URI string to register
    ///
    /// Returns:
    ///     int: The unique ID assigned to this URI
    ///
    /// Example:
    ///     >>> id = register.register_uri("http://example.org")
    ///     >>> print(f"Registered with ID: {id}")
    fn register_uri(&self, uri: String) -> PyResult<u64> {
        self.rt
            .block_on(self.inner.register_uri(&uri))
            .map_err(|e| PyRuntimeError::new_err(format!("Registration failed: {}", e)))
    }

    /// Register multiple URIs in batch and return their IDs (synchronous, order preserved)
    ///
    /// Args:
    ///     uris: List of URI strings to register
    ///
    /// Returns:
    ///     list[int]: List of IDs in the same order as input URIs
    ///
    /// Example:
    ///     >>> uris = ["http://example.org/1", "http://example.org/2"]
    ///     >>> ids = register.register_uri_batch(uris)
    fn register_uri_batch(&self, uris: Vec<String>) -> PyResult<Vec<u64>> {
        self.rt
            .block_on(self.inner.register_uri_batch(&uris))
            .map_err(|e| PyRuntimeError::new_err(format!("Batch registration failed: {}", e)))
    }

    /// Register multiple URIs in batch and return a dict of URI-to-ID mappings (synchronous)
    ///
    /// Args:
    ///     uris: List of URI strings to register
    ///
    /// Returns:
    ///     dict[str, int]: Dictionary mapping each URI to its assigned ID
    ///
    /// Example:
    ///     >>> uris = ["http://example.org/1", "http://example.org/2"]
    ///     >>> mapping = register.register_uri_batch_hashmap(uris)
    fn register_uri_batch_hashmap(&self, uris: Vec<String>) -> PyResult<HashMap<String, u64>> {
        self.rt
            .block_on(self.inner.register_uri_batch_hashmap(&uris))
            .map_err(|e| {
                PyRuntimeError::new_err(format!("Batch hashmap registration failed: {}", e))
            })
    }

    /// Get statistics about the URI register (synchronous)
    ///
    /// Returns:
    ///     dict: Dictionary with 'total_uris' and 'size_bytes' keys
    ///
    /// Example:
    ///     >>> stats = register.stats()
    ///     >>> print(f"Total URIs: {stats['total_uris']}")
    fn stats(&self) -> PyResult<HashMap<&'static str, u64>> {
        let stats = self
            .rt
            .block_on(self.inner.stats())
            .map_err(|e| PyRuntimeError::new_err(format!("Failed to get stats: {}", e)))?;

        let mut result = HashMap::new();
        result.insert("total_uris", stats.total_uris);
        result.insert("size_bytes", stats.size_bytes);
        Ok(result)
    }

    // ==================== Asynchronous methods ====================

    /// Register a single URI and return its ID (asynchronous)
    ///
    /// Args:
    ///     uri: The URI string to register
    ///
    /// Returns:
    ///     int: The unique ID assigned to this URI
    ///
    /// Example:
    ///     >>> id = await register.register_uri_async("http://example.org")
    ///     >>> print(f"Registered with ID: {id}")
    fn register_uri_async<'py>(&self, py: Python<'py>, uri: String) -> PyResult<Bound<'py, PyAny>> {
        let inner = self.inner.clone_inner();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let id = inner
                .register_uri(&uri)
                .await
                .map_err(|e| PyRuntimeError::new_err(format!("Registration failed: {}", e)))?;

            Ok(id)
        })
    }

    /// Register multiple URIs in batch and return their IDs (asynchronous, order preserved)
    ///
    /// Args:
    ///     uris: List of URI strings to register
    ///
    /// Returns:
    ///     list[int]: List of IDs in the same order as input URIs
    ///
    /// Example:
    ///     >>> uris = ["http://example.org/1", "http://example.org/2"]
    ///     >>> ids = await register.register_uri_batch_async(uris)
    fn register_uri_batch_async<'py>(
        &self,
        py: Python<'py>,
        uris: Vec<String>,
    ) -> PyResult<Bound<'py, PyAny>> {
        let inner = self.inner.clone_inner();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let ids = inner.register_uri_batch(&uris).await.map_err(|e| {
                PyRuntimeError::new_err(format!("Batch registration failed: {}", e))
            })?;

            Ok(ids)
        })
    }

    /// Register multiple URIs in batch and return a dict of URI-to-ID mappings (asynchronous)
    ///
    /// Args:
    ///     uris: List of URI strings to register
    ///
    /// Returns:
    ///     dict[str, int]: Dictionary mapping each URI to its assigned ID
    ///
    /// Example:
    ///     >>> uris = ["http://example.org/1", "http://example.org/2"]
    ///     >>> mapping = await register.register_uri_batch_hashmap_async(uris)
    fn register_uri_batch_hashmap_async<'py>(
        &self,
        py: Python<'py>,
        uris: Vec<String>,
    ) -> PyResult<Bound<'py, PyAny>> {
        let inner = self.inner.clone_inner();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let map = inner.register_uri_batch_hashmap(&uris).await.map_err(|e| {
                PyRuntimeError::new_err(format!("Batch hashmap registration failed: {}", e))
            })?;

            Ok(map)
        })
    }

    /// Get statistics about the URI register (asynchronous)
    ///
    /// Returns:
    ///     dict: Dictionary with 'total_uris' and 'size_bytes' keys
    ///
    /// Example:
    ///     >>> stats = await register.stats_async()
    ///     >>> print(f"Total URIs: {stats['total_uris']}")
    fn stats_async<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        let inner = self.inner.clone_inner();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let stats = inner
                .stats()
                .await
                .map_err(|e| PyRuntimeError::new_err(format!("Failed to get stats: {}", e)))?;

            let mut result = HashMap::new();
            result.insert("total_uris", stats.total_uris);
            result.insert("size_bytes", stats.size_bytes);

            Ok(result)
        })
    }

    fn __repr__(&self) -> String {
        "UriRegister(connected)".to_string()
    }
}

/// Python module initialization
#[pymodule]
fn _uri_register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    // Initialize logging bridge: tracing -> log -> Python logging
    // This allows Python apps to see Rust tracing events in their logging output
    // Use try_init() variants to avoid panics if loggers are already initialized
    tracing_log::LogTracer::init().ok(); // Convert tracing events to log events
    let _ = pyo3_log::try_init(); // Bridge log events to Python's logging module (ignore if already set)

    m.add_class::<PyUriRegister>()?;

    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
    m.add(
        "__doc__",
        "URI Register - A high-performance PostgreSQL-backed URI to ID mapping service",
    )?;

    Ok(())
}