formatparse-pyo3 0.8.2

use crate::error;
use crate::parser::findall_iter::FindallIter;
use crate::parser::format_parser::{CompiledFields, FormatParser};
use crate::pattern_cache::get_or_create_parser;
use fancy_regex::Regex;
use formatparse_core::count_capturing_groups;
use formatparse_core::parser::validate_input_length;
use formatparse_core::{FieldSpec, FieldType};
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::{PyAnyMethods, PyDict, PyList, PyString, PyTuple};
use pyo3::IntoPyObjectExt;
use std::collections::HashMap;

fn merge_call_extra_types(
    py: Python<'_>,
    stored: &Option<HashMap<String, Py<PyAny>>>,
    provided: Option<HashMap<String, Py<PyAny>>>,
) -> PyResult<Option<HashMap<String, Py<PyAny>>>> {
    let mut merged = if let Some(ref stored) = stored {
        stored
            .iter()
            .map(|(k, v)| (k.clone(), v.clone_ref(py)))
            .collect()
    } else {
        HashMap::new()
    };
    if let Some(provided) = provided {
        for (k, v) in provided {
            merged.insert(k.clone(), v.clone_ref(py));
        }
    }
    Ok(Some(merged))
}

fn field_type_validation_tag(spec: &FieldSpec) -> String {
    match &spec.field_type {
        FieldType::String => "s".to_string(),
        FieldType::Integer => "d".to_string(),
        FieldType::Float => "f".to_string(),
        FieldType::Boolean => "b".to_string(),
        FieldType::Letters => "l".to_string(),
        FieldType::Word => "w".to_string(),
        FieldType::NonLetters => "W".to_string(),
        FieldType::NonWhitespace => "S".to_string(),
        FieldType::NonDigits => "D".to_string(),
        FieldType::Custom(name) => name.clone(),
        FieldType::Nested => "s".to_string(),
        FieldType::Multiline | FieldType::IndentBlock => "s".to_string(),
        FieldType::BracedContent => "s".to_string(),
        FieldType::NumberWithThousands => "n".to_string(),
        FieldType::Scientific => "e".to_string(),
        FieldType::GeneralNumber => "g".to_string(),
        FieldType::Percentage => "%".to_string(),
        FieldType::DateTimeISO => "ti".to_string(),
        FieldType::DateTimeRFC2822 => "te".to_string(),
        FieldType::DateTimeGlobal => "tg".to_string(),
        FieldType::DateTimeUS => "ta".to_string(),
        FieldType::DateTimeCtime => "tc".to_string(),
        FieldType::DateTimeHTTP => "th".to_string(),
        FieldType::DateTimeTime => "tt".to_string(),
        FieldType::DateTimeSystem => "ts".to_string(),
        FieldType::DateTimeStrftime => "%".to_string(),
    }
}

#[pymethods]
impl FormatParser {
    #[new]
    #[pyo3(signature = (pattern=None, extra_types=None))]
    fn new_py(
        pattern: Option<&str>,
        extra_types: Option<HashMap<String, Py<PyAny>>>,
    ) -> PyResult<Self> {
        match pattern {
            Some(p) => Self::new_with_extra_types(p, extra_types),
            None => {
                // Create a dummy instance for unpickling - __setstate__ will initialize it properly
                // We need to create a valid but minimal instance
                let empty_regex =
                    Regex::new("^$").map_err(|e| crate::error::regex_error(&e.to_string()))?;
                Ok(Self {
                    pattern: String::new(),
                    regex: empty_regex.clone(),
                    regex_str: String::new(),
                    regex_case_insensitive: None,
                    search_regex: empty_regex.clone(),
                    search_regex_case_insensitive: None,
                    fields: CompiledFields {
                        field_specs: Vec::new(),
                        field_names: Vec::new(),
                        normalized_names: Vec::new(),
                        custom_type_groups: Vec::new(),
                        has_nested_dict_fields: Vec::new(),
                        nested_parsers: Vec::new(),
                        field_count: 0,
                    },
                    name_mapping: HashMap::new(),
                    stored_extra_types: None,
                    allows_empty_default_string_match: false,
                })
            }
        }
    }

    /// Parse a string using this compiled pattern.
    ///
    /// Merges ``extra_types`` from compile time with any ``extra_types`` passed here
    /// (call-time wins on duplicate keys).
    #[pyo3(signature = (string, case_sensitive=false, extra_types=None, evaluate_result=true))]
    fn parse(
        &self,
        string: &str,
        case_sensitive: bool,
        extra_types: Option<HashMap<String, Py<PyAny>>>,
        evaluate_result: bool,
    ) -> PyResult<Option<Py<PyAny>>> {
        // Validate input length
        validate_input_length(string).map_err(PyValueError::new_err)?;

        // Check for null bytes
        if string.contains('\0') {
            return Err(PyValueError::new_err("Input string contains null byte"));
        }
        let merged_extra_types = Python::attach(|py| {
            merge_call_extra_types(py, &self.stored_extra_types, extra_types)
        })?;
        let parser = Python::attach(|py| -> PyResult<std::sync::Arc<FormatParser>> {
            let cache_et = merged_extra_types.as_ref().map(|m| {
                m.iter()
                    .map(|(k, v)| (k.clone(), v.clone_ref(py)))
                    .collect()
            });
            get_or_create_parser(&self.pattern, cache_et)
        })?;
        parser.parse_internal(
            string,
            case_sensitive,
            merged_extra_types.as_ref(),
            evaluate_result,
        )
    }

    /// Get the list of named field names (returns normalized names for compatibility)
    #[getter]
    fn named_fields(&self) -> Vec<String> {
        // Return normalized names (without hyphens/dots) for compatibility with original parse library
        self.fields
            .normalized_names
            .iter()
            .filter_map(|n| n.clone())
            .collect()
    }

    /// Per-field layout for bidirectional validation (name, type tag, width, precision).
    #[getter]
    fn field_constraints(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let list = PyList::empty(py);
        for (i, spec) in self.fields.field_specs.iter().enumerate() {
            let dict = PyDict::new(py);
            match self.fields.field_names.get(i).and_then(|n| n.as_deref()) {
                Some(name) => dict.set_item("name", name)?,
                None => dict.set_item("name", py.None())?,
            }
            dict.set_item("type", field_type_validation_tag(spec))?;
            match spec.width {
                Some(w) => dict.set_item("width", w as i64)?,
                None => dict.set_item("width", py.None())?,
            }
            match spec.precision {
                Some(p) => dict.set_item("precision", p as i64)?,
                None => dict.set_item("precision", py.None())?,
            }
            list.append(dict)?;
        }
        list.into_py_any(py)
    }

    /// Raw regex body for this pattern (no ``^``/``$``, no ``(?s)`` prefix).
    ///
    /// Intended for **composition** (GitHub issue #7): embed this string as a custom
    /// type’s ``pattern`` when building a parent pattern via ``extra_types``. Do not use
    /// :meth:`_expression` for that purpose; it applies display-oriented transforms that are
    /// not guaranteed to be valid regex fragments.
    #[getter]
    fn regex_subpattern(&self) -> String {
        self.regex_str.clone()
    }

    /// Number of capturing groups in :meth:`regex_subpattern`.
    ///
    /// Use as ``regex_group_count`` on a converter object when composing (see
    /// :func:`formatparse.composed_type`).
    #[getter]
    fn regex_capturing_group_count(&self) -> usize {
        count_capturing_groups(&self.regex_str)
    }

    /// Get the internal regex expression string (for testing)
    /// Returns a canonical format with literal spaces instead of \s+ for compatibility
    #[getter]
    fn _expression(&self) -> String {
        let mut result = self.regex_str.clone();

        // Replace \s+ between capturing groups with literal spaces for canonical format
        // This matches the original parse library's _expression format
        result = result.replace(r")\s+(", ") (");
        // Also replace )\s*( with ) ( for backward compatibility
        result = result.replace(r")\s*(", ") (");

        // Simplify float patterns to match expected format
        // Our pattern: ([+-]?(?:\d+\.\d+|\.\d+|\d+\.)(?:[eE][+-]?\d+)?)
        // Expected: ([-+ ]?\d*\.\d+)
        // Replace the complex float pattern with the simpler one
        result = result.replace(
            r"([+-]?(?:\d+\.\d+|\.\d+|\d+\.)(?:[eE][+-]?\d+)?)",
            r"([-+ ]?\d*\.\d+)",
        );

        // For alignment patterns like {:>} that produce "( *(.+?))", we need to unwrap
        // the outer capturing group to get " *(.+?)" (no outer wrapper)
        // Only do this for patterns that start with "(" and end with ")" and contain nested groups
        if result.starts_with("(") && result.ends_with(")") {
            let inner = &result[1..result.len() - 1];
            // Check if inner already starts with a space and contains a capturing group
            if inner.starts_with(" *(") && inner.ends_with(")") {
                // This is a simple wrapper, unwrap it
                result = inner.to_string();
            }
        }

        result
    }

    /// Get the format object for formatting values into the pattern
    #[getter]
    fn format(&self) -> Format {
        Format {
            pattern: self.pattern.clone(),
        }
    }

    /// Search for the pattern in a string.
    ///
    /// Merges ``extra_types`` from compile time with any ``extra_types`` passed here
    /// (call-time wins on duplicate keys), same as :meth:`parse`.
    #[pyo3(signature = (string, case_sensitive=true, extra_types=None, evaluate_result=true))]
    fn search(
        &self,
        string: &str,
        case_sensitive: bool,
        extra_types: Option<HashMap<String, Py<PyAny>>>,
        evaluate_result: bool,
    ) -> PyResult<Option<Py<PyAny>>> {
        // Validate input length
        validate_input_length(string).map_err(PyValueError::new_err)?;

        // Check for null bytes
        if string.contains('\0') {
            return Err(PyValueError::new_err("Input string contains null byte"));
        }

        let merged_extra_types = Python::attach(|py| {
            merge_call_extra_types(py, &self.stored_extra_types, extra_types)
        })?;
        let parser = Python::attach(|py| -> PyResult<std::sync::Arc<FormatParser>> {
            let cache_et = merged_extra_types.as_ref().map(|m| {
                m.iter()
                    .map(|(k, v)| (k.clone(), v.clone_ref(py)))
                    .collect()
            });
            get_or_create_parser(&self.pattern, cache_et)
        })?;
        parser.search_pattern(string, case_sensitive, merged_extra_types, evaluate_result)
    }

    /// Yield non-overlapping matches from ``string`` one at a time.
    ///
    /// Same matching rules as :func:`findall`, but each ``__next__`` converts at most one
    /// match, lowering peak memory when you stream results. This does **not** read
    /// arbitrary file chunks with backtracking; pair with line-based file iteration only
    /// when matches cannot span line breaks (see GitHub issue #13).
    #[pyo3(signature = (string, case_sensitive=false, extra_types=None, evaluate_result=true, max_matches=None))]
    fn findall_iter(
        &self,
        py: Python<'_>,
        string: &str,
        case_sensitive: bool,
        extra_types: Option<HashMap<String, Py<PyAny>>>,
        evaluate_result: bool,
        max_matches: Option<usize>,
    ) -> PyResult<Py<FindallIter>> {
        validate_input_length(string).map_err(PyValueError::new_err)?;

        if string.contains('\0') {
            return Err(PyValueError::new_err("Input string contains null byte"));
        }

        let merged_extra_types = Python::attach(|py| {
            merge_call_extra_types(py, &self.stored_extra_types, extra_types)
        })?;
        let merged_map = merged_extra_types.unwrap_or_default();
        let parser = Python::attach(|py| -> PyResult<std::sync::Arc<FormatParser>> {
            let cache_et = if merged_map.is_empty() {
                None
            } else {
                Some(
                    merged_map
                        .iter()
                        .map(|(k, v)| (k.clone(), v.clone_ref(py)))
                        .collect(),
                )
            };
            get_or_create_parser(&self.pattern, cache_et)
        })?;
        Py::new(
            py,
            FindallIter::new(
                parser,
                string.to_string(),
                case_sensitive,
                evaluate_result,
                merged_map,
                max_matches,
            ),
        )
    }

    /// Pickle support: rebuild with ``compile(pattern)`` only (no ``extra_types``).
    ///
    /// Custom type converters cannot be serialized reliably; use
    /// ``compile(pattern, extra_types=...)`` after unpickling if you need them.
    fn __reduce__(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let m = py.import("_formatparse")?;
        let compile_fn = m.getattr("compile")?;
        let args = PyTuple::new(py, [&self.pattern])?;
        PyTuple::new(py, [compile_fn.as_any(), args.as_any()])?.into_py_any(py)
    }

    /// Pickle state: pattern string only (see class doc for ``extra_types``).
    fn __getstate__(&self, py: Python) -> PyResult<Py<PyAny>> {
        use pyo3::types::PyDict;
        let state = PyDict::new(py);
        state.set_item("pattern", &self.pattern)?;
        state.into_py_any(py)
    }

    /// Restore from pickle state (pattern only; ``extra_types`` are not recovered).
    fn __setstate__(&mut self, _py: Python, state: &Bound<'_, PyAny>) -> PyResult<()> {
        use pyo3::types::PyDict;
        let dict = state.cast::<PyDict>()?;
        let pattern: String = dict
            .get_item("pattern")?
            .ok_or_else(|| error::missing_field_error("pattern"))?
            .extract()?;

        // Reconstruct the parser from the pattern
        let reconstructed = Self::new_with_extra_types(&pattern, None)?;

        // Copy all fields from reconstructed parser
        self.pattern = reconstructed.pattern;
        self.regex_str = reconstructed.regex_str;
        self.regex = reconstructed.regex;
        self.regex_case_insensitive = reconstructed.regex_case_insensitive;
        self.search_regex = reconstructed.search_regex;
        self.search_regex_case_insensitive = reconstructed.search_regex_case_insensitive;
        self.fields = reconstructed.fields;
        self.name_mapping = reconstructed.name_mapping;
        self.stored_extra_types = reconstructed.stored_extra_types;
        self.allows_empty_default_string_match = reconstructed.allows_empty_default_string_match;
        Ok(())
    }
}

/// Format object that formats values into a pattern string
#[pyclass]
pub struct Format {
    pattern: String,
}

#[pymethods]
impl Format {
    /// Format values into the pattern string using Python's format() method
    fn format(&self, py: Python, args: &Bound<'_, PyAny>) -> PyResult<String> {
        // Use Python's string format method to format values into the pattern
        let pattern_obj = PyString::new(py, &self.pattern);
        let format_method = pattern_obj.getattr("format")?;

        // Call format with the args (can be a single value, tuple, or *args)
        let result = if let Ok(tuple) = args.cast::<PyTuple>() {
            format_method.call1(tuple)?
        } else {
            // Single argument
            format_method.call1((args,))?
        };
        result.extract()
    }
}