prefix-register 0.2.2

// 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 prefix-register using PyO3

use crate::{PrefixRegistry, RetryConfig};
use pyo3::exceptions::{PyRuntimeError, PyValueError};
use pyo3::prelude::*;
use std::collections::HashMap;
use std::time::Duration;

/// Python wrapper for PrefixRegistry
#[pyclass(name = "PrefixRegistry")]
struct PyPrefixRegistry {
    inner: PrefixRegistry,
}

#[pymethods]
impl PyPrefixRegistry {
    /// Create a new prefix registry connected to PostgreSQL
    ///
    /// Args:
    ///     database_url: PostgreSQL connection string (e.g., "postgres://user:password@host:port/database")
    ///     max_connections: Maximum number of connections in the pool (recommended: 5-20)
    ///
    /// Returns:
    ///     PrefixRegistry: A new prefix registry instance
    ///
    /// Example:
    ///     >>> registry = await PrefixRegistry.new("postgres://localhost/mydb", 10)
    #[staticmethod]
    #[allow(clippy::new_ret_no_self)]
    #[pyo3(signature = (database_url, max_connections))]
    fn new<'p>(
        py: Python<'p>,
        database_url: String,
        max_connections: u32,
    ) -> PyResult<Bound<'p, PyAny>> {
        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",
            ));
        }

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let registry = PrefixRegistry::new(&database_url, max_connections as usize)
                .await
                .map_err(|e| PyRuntimeError::new_err(format!("Failed to connect: {}", e)))?;

            Ok(PyPrefixRegistry { inner: registry })
        })
    }

    /// Create a new prefix registry with retry logic for transient failures
    ///
    /// Args:
    ///     database_url: PostgreSQL connection string
    ///     max_connections: Maximum number of connections in the pool
    ///     max_retries: Maximum number of retry attempts (default: 5)
    ///     initial_delay_ms: Initial delay in milliseconds before first retry (default: 1000)
    ///     max_delay_ms: Maximum delay in milliseconds between retries (default: 30000)
    ///
    /// Returns:
    ///     PrefixRegistry: A new prefix registry instance
    ///
    /// Example:
    ///     >>> registry = await PrefixRegistry.new_with_retry(
    ///     ...     "postgres://localhost/mydb", 10, 5, 1000, 30000
    ///     ... )
    #[staticmethod]
    #[pyo3(signature = (database_url, max_connections, max_retries=5, initial_delay_ms=1000, max_delay_ms=30000))]
    fn new_with_retry<'p>(
        py: Python<'p>,
        database_url: String,
        max_connections: u32,
        max_retries: u32,
        initial_delay_ms: u64,
        max_delay_ms: u64,
    ) -> PyResult<Bound<'p, PyAny>> {
        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",
            ));
        }

        let retry_config = RetryConfig::new(
            max_retries,
            Duration::from_millis(initial_delay_ms),
            Duration::from_millis(max_delay_ms),
        );

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let registry = PrefixRegistry::new_with_retry(
                &database_url,
                max_connections as usize,
                retry_config,
            )
            .await
            .map_err(|e| PyRuntimeError::new_err(format!("Failed to connect: {}", e)))?;

            Ok(PyPrefixRegistry { inner: registry })
        })
    }

    /// Store a new prefix if the URI doesn't already have one
    ///
    /// Args:
    ///     prefix: The namespace prefix (e.g., "foaf", "rdf")
    ///     uri: The full namespace URI (e.g., "http://xmlns.com/foaf/0.1/")
    ///
    /// Returns:
    ///     bool: True if the prefix was stored, False if the URI already had a prefix
    ///
    /// Example:
    ///     >>> stored = await registry.store_prefix_if_new("foaf", "http://xmlns.com/foaf/0.1/")
    ///     >>> if stored:
    ///     ...     print("New prefix stored")
    fn store_prefix_if_new<'p>(
        &self,
        py: Python<'p>,
        prefix: String,
        uri: String,
    ) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();

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

            Ok(stored)
        })
    }

    /// Store multiple prefixes in batch
    ///
    /// Args:
    ///     prefixes: List of (prefix, uri) tuples to store
    ///
    /// Returns:
    ///     dict: Dictionary with 'stored' and 'skipped' counts
    ///
    /// Example:
    ///     >>> prefixes = [("foaf", "http://xmlns.com/foaf/0.1/"), ("rdf", "http://www.w3.org/1999/02/22-rdf-syntax-ns#")]
    ///     >>> result = await registry.store_prefixes_if_new(prefixes)
    ///     >>> print(f"Stored {result['stored']}, skipped {result['skipped']}")
    fn store_prefixes_if_new<'p>(
        &self,
        py: Python<'p>,
        prefixes: Vec<(String, String)>,
    ) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            // Convert owned strings to references for the API
            let prefix_refs: Vec<(&str, &str)> = prefixes
                .iter()
                .map(|(p, u)| (p.as_str(), u.as_str()))
                .collect();

            let result = inner
                .store_prefixes_if_new(prefix_refs)
                .await
                .map_err(|e| PyRuntimeError::new_err(format!("Batch store failed: {}", e)))?;

            let mut dict = HashMap::new();
            dict.insert("stored", result.stored);
            dict.insert("skipped", result.skipped);

            Ok(dict)
        })
    }

    /// Get the URI for a given prefix
    ///
    /// Args:
    ///     prefix: The namespace prefix (e.g., "foaf")
    ///
    /// Returns:
    ///     Optional[str]: The URI if the prefix is known, None otherwise
    ///
    /// Example:
    ///     >>> uri = await registry.get_uri_for_prefix("foaf")
    ///     >>> if uri:
    ///     ...     print(f"foaf = {uri}")
    fn get_uri_for_prefix<'p>(&self, py: Python<'p>, prefix: String) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();

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

            Ok(uri)
        })
    }

    /// Get the prefix for a given URI
    ///
    /// Args:
    ///     uri: The full namespace URI
    ///
    /// Returns:
    ///     Optional[str]: The prefix if the URI is registered, None otherwise
    ///
    /// Example:
    ///     >>> prefix = await registry.get_prefix_for_uri("http://xmlns.com/foaf/0.1/")
    ///     >>> if prefix:
    ///     ...     print(f"URI has prefix: {prefix}")
    fn get_prefix_for_uri<'p>(&self, py: Python<'p>, uri: String) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();

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

            Ok(prefix)
        })
    }

    /// Expand a CURIE to a full URI
    ///
    /// Args:
    ///     prefix: The namespace prefix (e.g., "foaf")
    ///     local_name: The local part (e.g., "Person")
    ///
    /// Returns:
    ///     Optional[str]: The full URI (e.g., "http://xmlns.com/foaf/0.1/Person") or None if prefix unknown
    ///
    /// Example:
    ///     >>> uri = await registry.expand_curie("foaf", "Person")
    ///     >>> if uri:
    ///     ...     print(f"foaf:Person = {uri}")
    fn expand_curie<'p>(
        &self,
        py: Python<'p>,
        prefix: String,
        local_name: String,
    ) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();

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

            Ok(uri)
        })
    }

    /// Get all registered prefixes
    ///
    /// Returns:
    ///     dict[str, str]: Dictionary mapping prefixes to URIs
    ///
    /// Example:
    ///     >>> prefixes = await registry.get_all_prefixes()
    ///     >>> for prefix, uri in prefixes.items():
    ///     ...     print(f"{prefix}: {uri}")
    fn get_all_prefixes<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let prefixes = inner.get_all_prefixes().await;
            Ok(prefixes)
        })
    }

    /// Get the number of registered prefixes
    ///
    /// Returns:
    ///     int: The number of registered prefixes
    ///
    /// Example:
    ///     >>> count = await registry.prefix_count()
    ///     >>> print(f"Registered prefixes: {count}")
    fn prefix_count<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let count = inner.prefix_count().await;
            Ok(count)
        })
    }

    /// Shorten a URI to a (prefix, local_name) tuple
    ///
    /// Uses longest-match semantics: if multiple namespaces match the URI,
    /// the longest one wins.
    ///
    /// Args:
    ///     uri: The full URI to shorten
    ///
    /// Returns:
    ///     Optional[tuple[str, str]]: A (prefix, local_name) tuple if a match is found, None otherwise
    ///
    /// Example:
    ///     >>> result = await registry.shorten_uri("http://xmlns.com/foaf/0.1/Person")
    ///     >>> if result:
    ///     ...     prefix, local = result
    ///     ...     print(f"{prefix}:{local}")
    fn shorten_uri<'p>(&self, py: Python<'p>, uri: String) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();

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

            Ok(result)
        })
    }

    /// Shorten a URI to a CURIE string, or return the original URI if no match
    ///
    /// Args:
    ///     uri: The full URI to shorten
    ///
    /// Returns:
    ///     str: A CURIE string like "prefix:local", or the original URI if no namespace matches
    ///
    /// Example:
    ///     >>> result = await registry.shorten_uri_or_full("http://xmlns.com/foaf/0.1/Person")
    ///     >>> print(result)  # "foaf:Person" or the full URI
    fn shorten_uri_or_full<'p>(&self, py: Python<'p>, uri: String) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();

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

            Ok(result)
        })
    }

    /// Shorten multiple URIs in batch
    ///
    /// Args:
    ///     uris: List of URIs to shorten
    ///
    /// Returns:
    ///     list[tuple[str, str] | None]: List with (prefix, local_name) for matched URIs, None for unmatched
    ///
    /// Example:
    ///     >>> uris = ["http://xmlns.com/foaf/0.1/Person", "http://unknown.org/thing"]
    ///     >>> results = await registry.shorten_uri_batch(uris)
    ///     >>> for result in results:
    ///     ...     if result:
    ///     ...         print(f"{result[0]}:{result[1]}")
    ///     ...     else:
    ///     ...         print("No match")
    fn shorten_uri_batch<'p>(
        &self,
        py: Python<'p>,
        uris: Vec<String>,
    ) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let uri_refs: Vec<&str> = uris.iter().map(|s| s.as_str()).collect();

            let results = inner
                .shorten_uri_batch(uri_refs)
                .await
                .map_err(|e| PyRuntimeError::new_err(format!("Batch shorten failed: {}", e)))?;

            Ok(results)
        })
    }

    /// Expand multiple CURIEs in batch
    ///
    /// Args:
    ///     curies: List of (prefix, local_name) tuples
    ///
    /// Returns:
    ///     list[str | None]: List with expanded URIs for known prefixes, None for unknown
    ///
    /// Example:
    ///     >>> curies = [("foaf", "Person"), ("unknown", "Thing")]
    ///     >>> results = await registry.expand_curie_batch(curies)
    ///     >>> for result in results:
    ///     ...     if result:
    ///     ...         print(result)
    ///     ...     else:
    ///     ...         print("Unknown prefix")
    fn expand_curie_batch<'p>(
        &self,
        py: Python<'p>,
        curies: Vec<(String, String)>,
    ) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let curie_refs: Vec<(&str, &str)> = curies
                .iter()
                .map(|(p, l)| (p.as_str(), l.as_str()))
                .collect();

            let results = inner
                .expand_curie_batch(curie_refs)
                .await
                .map_err(|e| PyRuntimeError::new_err(format!("Batch expand failed: {}", e)))?;

            Ok(results)
        })
    }

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

/// Python module initialization
#[pymodule]
fn _prefix_register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyPrefixRegistry>()?;

    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
    m.add(
        "__doc__",
        "Prefix Register - A PostgreSQL-backed namespace prefix registry for CURIE expansion",
    )?;

    Ok(())
}