helios-persistence 0.1.38

//! SearchParameter Value Extractor.
//!
//! Uses FHIRPath expressions to extract searchable values from FHIR resources.

use std::collections::HashMap;
use std::sync::Arc;

use helios_fhirpath::EvaluationContext;
use helios_fhirpath_support::EvaluationResult;
use parking_lot::RwLock;
use rust_decimal::Decimal;
use serde::{Deserialize, Serialize};
use serde_json::Value;

use crate::types::SearchParamType;

use super::converters::{IndexValue, ValueConverter};
use super::errors::ExtractionError;
use super::registry::{SearchParameterDefinition, SearchParameterRegistry};

/// A value extracted from a resource for indexing.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExtractedValue {
    /// The parameter name (e.g., "name", "identifier").
    pub param_name: String,

    /// The parameter URL.
    pub param_url: String,

    /// The parameter type.
    pub param_type: SearchParamType,

    /// The extracted and converted value.
    pub value: IndexValue,

    /// Composite group ID (for composite parameters).
    /// Values with the same group ID are part of the same composite match.
    pub composite_group: Option<u32>,
}

impl ExtractedValue {
    /// Creates a new extracted value.
    pub fn new(
        param_name: impl Into<String>,
        param_url: impl Into<String>,
        param_type: SearchParamType,
        value: IndexValue,
    ) -> Self {
        Self {
            param_name: param_name.into(),
            param_url: param_url.into(),
            param_type,
            value,
            composite_group: None,
        }
    }

    /// Sets the composite group ID.
    pub fn with_composite_group(mut self, group: u32) -> Self {
        self.composite_group = Some(group);
        self
    }
}

/// Extracts searchable values from FHIR resources using FHIRPath.
pub struct SearchParameterExtractor {
    registry: Arc<RwLock<SearchParameterRegistry>>,
}

impl SearchParameterExtractor {
    /// Creates a new extractor with the given registry.
    pub fn new(registry: Arc<RwLock<SearchParameterRegistry>>) -> Self {
        Self { registry }
    }

    /// Extracts all searchable values from a resource.
    ///
    /// Returns values for all active search parameters that apply to this resource type.
    pub fn extract(
        &self,
        resource: &Value,
        resource_type: &str,
    ) -> Result<Vec<ExtractedValue>, ExtractionError> {
        // Validate resource
        let obj = resource
            .as_object()
            .ok_or_else(|| ExtractionError::InvalidResource {
                message: "Resource must be a JSON object".to_string(),
            })?;

        // Verify resource type
        if let Some(rt) = obj.get("resourceType").and_then(|v| v.as_str()) {
            if rt != resource_type {
                return Err(ExtractionError::InvalidResource {
                    message: format!(
                        "Resource type mismatch: expected {}, got {}",
                        resource_type, rt
                    ),
                });
            }
        }

        let mut results = Vec::new();

        // Get active parameters for this resource type
        let params = {
            let registry = self.registry.read();
            registry.get_active_params(resource_type)
        };

        for param in &params {
            match self.extract_for_param(resource, param) {
                Ok(values) => results.extend(values),
                Err(e) => {
                    // Log the error but continue with other parameters
                    tracing::warn!(
                        "Failed to extract values for parameter '{}': {}",
                        param.code,
                        e
                    );
                }
            }
        }

        // Also extract common Resource-level parameters
        let common_params = {
            let registry = self.registry.read();
            registry.get_active_params("Resource")
        };

        for param in &common_params {
            if !params.iter().any(|p| p.code == param.code) {
                match self.extract_for_param(resource, param) {
                    Ok(values) => results.extend(values),
                    Err(e) => {
                        tracing::warn!(
                            "Failed to extract values for common parameter '{}': {}",
                            param.code,
                            e
                        );
                    }
                }
            }
        }

        Ok(results)
    }

    /// Extracts values for a specific parameter from a resource.
    pub fn extract_for_param(
        &self,
        resource: &Value,
        param: &SearchParameterDefinition,
    ) -> Result<Vec<ExtractedValue>, ExtractionError> {
        if param.expression.is_empty() {
            return Ok(Vec::new());
        }

        // Get the resource type from the resource
        let resource_type = resource
            .get("resourceType")
            .and_then(|v| v.as_str())
            .unwrap_or("");

        // Filter the expression to only include parts relevant to this resource type
        let filtered_expr = self.filter_expression_for_resource(&param.expression, resource_type);

        if filtered_expr.is_empty() {
            return Ok(Vec::new());
        }

        // Evaluate the filtered FHIRPath expression using the actual evaluator
        let values = self.evaluate_fhirpath(resource, &filtered_expr)?;

        let mut results = Vec::new();
        for value in values {
            let converted = ValueConverter::convert(&value, param.param_type, &param.code)?;
            for idx_value in converted {
                results.push(ExtractedValue::new(
                    &param.code,
                    &param.url,
                    param.param_type,
                    idx_value,
                ));
            }
        }

        Ok(results)
    }

    /// Filters a FHIRPath expression to only include parts relevant to a specific resource type.
    ///
    /// Many FHIR SearchParameters have expressions that span multiple resource types, joined
    /// with `|` (union). For example, the `patient` parameter has:
    /// `AllergyIntolerance.patient | CarePlan.subject.where(resolve() is Patient) | ...`
    ///
    /// This method extracts only the parts that start with the given resource type and
    /// simplifies common patterns that use `resolve()`.
    fn filter_expression_for_resource(&self, expression: &str, resource_type: &str) -> String {
        // Split by | and filter to parts starting with our resource type
        let parts: Vec<String> = expression
            .split('|')
            .map(|p| p.trim())
            .filter(|p| {
                // Check if this part starts with our resource type
                p.starts_with(resource_type)
                    && (p.len() == resource_type.len()
                        || p.chars().nth(resource_type.len()) == Some('.'))
            })
            .map(|p| self.simplify_resolve_pattern(p))
            .collect();

        if parts.is_empty() {
            // If no parts match, return the original expression
            // This handles expressions that don't use ResourceType prefix
            expression.to_string()
        } else {
            // Join the filtered parts back with |
            parts.join(" | ")
        }
    }

    /// Simplifies common `.where(resolve() is ResourceType)` patterns.
    ///
    /// In FHIR SearchParameters, patterns like `subject.where(resolve() is Patient)`
    /// are used to filter references by target type. Since we're extracting references
    /// for indexing (not actually resolving them), we can safely strip this pattern
    /// and just extract the reference value.
    fn simplify_resolve_pattern(&self, expr: &str) -> String {
        // Pattern: .where(resolve() is SomeType)
        // We want to remove this suffix since we just need the reference value
        if let Some(where_pos) = expr.find(".where(resolve()") {
            // Find the matching closing paren
            let after_where = &expr[where_pos..];
            if after_where.rfind(')').is_some() {
                // Return everything before .where(...)
                return expr[..where_pos].to_string();
            }
        }
        expr.to_string()
    }

    /// Evaluates a FHIRPath expression against a resource using the helios-fhirpath evaluator.
    fn evaluate_fhirpath(
        &self,
        resource: &Value,
        expression: &str,
    ) -> Result<Vec<Value>, ExtractionError> {
        // Convert JSON to EvaluationResult and set up context
        let eval_result = json_to_evaluation_result(resource)?;

        // Create evaluation context with the resource as 'this'
        let mut context = EvaluationContext::new_empty_with_default_version();
        context.set_this(eval_result);

        // Evaluate the FHIRPath expression
        let result = helios_fhirpath::evaluate_expression(expression, &context).map_err(|e| {
            ExtractionError::FhirPathError {
                expression: expression.to_string(),
                message: e,
            }
        })?;

        // Convert EvaluationResult back to JSON values
        evaluation_result_to_json_values(&result)
    }
}

/// Converts a serde_json::Value to an EvaluationResult.
fn json_to_evaluation_result(value: &Value) -> Result<EvaluationResult, ExtractionError> {
    match value {
        Value::Null => Ok(EvaluationResult::Empty),
        Value::Bool(b) => Ok(EvaluationResult::boolean(*b)),
        Value::Number(n) => {
            if let Some(i) = n.as_i64() {
                Ok(EvaluationResult::integer(i))
            } else if let Some(f) = n.as_f64() {
                Ok(EvaluationResult::decimal(Decimal::try_from(f).map_err(
                    |e| ExtractionError::ConversionError {
                        message: format!("Invalid decimal: {}", e),
                    },
                )?))
            } else {
                Err(ExtractionError::ConversionError {
                    message: "Invalid number".to_string(),
                })
            }
        }
        Value::String(s) => Ok(EvaluationResult::string(s.clone())),
        Value::Array(arr) => {
            let results: Result<Vec<_>, _> = arr.iter().map(json_to_evaluation_result).collect();
            Ok(EvaluationResult::collection(results?))
        }
        Value::Object(obj) => {
            let mut map = HashMap::new();
            for (key, val) in obj {
                let eval_val = json_to_evaluation_result(val)?;
                map.insert(key.clone(), eval_val);
            }
            Ok(EvaluationResult::Object {
                map,
                type_info: None,
            })
        }
    }
}

/// Converts an EvaluationResult back to JSON values for the converter.
fn evaluation_result_to_json_values(
    result: &EvaluationResult,
) -> Result<Vec<Value>, ExtractionError> {
    match result {
        EvaluationResult::Empty => Ok(Vec::new()),
        EvaluationResult::Boolean(b, _) => Ok(vec![Value::Bool(*b)]),
        EvaluationResult::String(s, _) => Ok(vec![Value::String(s.clone())]),
        EvaluationResult::Integer(i, _) => Ok(vec![Value::Number((*i).into())]),
        EvaluationResult::Integer64(i, _) => Ok(vec![Value::Number((*i).into())]),
        EvaluationResult::Decimal(d, _) => {
            // Convert decimal to JSON number
            let f: f64 = (*d).try_into().unwrap_or(0.0);
            Ok(vec![Value::Number(
                serde_json::Number::from_f64(f).unwrap_or_else(|| serde_json::Number::from(0)),
            )])
        }
        EvaluationResult::Date(s, _) => Ok(vec![Value::String(s.clone())]),
        EvaluationResult::DateTime(s, _) => Ok(vec![Value::String(s.clone())]),
        EvaluationResult::Time(s, _) => Ok(vec![Value::String(s.clone())]),
        EvaluationResult::Quantity(value, unit, _) => {
            // Convert Quantity to JSON object
            let f: f64 = (*value).try_into().unwrap_or(0.0);
            Ok(vec![serde_json::json!({
                "value": f,
                "unit": unit
            })])
        }
        EvaluationResult::Collection { items, .. } => {
            let mut values = Vec::new();
            for item in items {
                values.extend(evaluation_result_to_json_values(item)?);
            }
            Ok(values)
        }
        EvaluationResult::Object { map, .. } => {
            // Convert object back to JSON
            let mut obj = serde_json::Map::new();
            for (key, val) in map {
                let json_vals = evaluation_result_to_json_values(val)?;
                // Check if the original value was a Collection - if so, preserve it as an array
                // even if it has only one element, since FHIR arrays should stay as arrays
                let is_collection = matches!(val, EvaluationResult::Collection { .. });
                if is_collection {
                    // Always preserve arrays as arrays
                    obj.insert(key.clone(), Value::Array(json_vals));
                } else if json_vals.len() == 1 {
                    obj.insert(key.clone(), json_vals.into_iter().next().unwrap());
                } else if !json_vals.is_empty() {
                    obj.insert(key.clone(), Value::Array(json_vals));
                }
            }
            Ok(vec![Value::Object(obj)])
        }
    }
}

impl std::fmt::Debug for SearchParameterExtractor {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SearchParameterExtractor").finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::search::loader::SearchParameterLoader;
    use helios_fhir::FhirVersion;
    use serde_json::json;
    use std::path::PathBuf;

    fn create_test_extractor() -> SearchParameterExtractor {
        let loader = SearchParameterLoader::new(FhirVersion::R4);
        let mut registry = SearchParameterRegistry::new();

        // Load minimal fallback
        if let Ok(params) = loader.load_embedded() {
            for param in params {
                let _ = registry.register(param);
            }
        }

        // Load spec file for full parameter support
        // CARGO_MANIFEST_DIR for this crate is crates/persistence
        // We need to go up two levels to reach the workspace root
        let data_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
            .parent()
            .and_then(|p| p.parent())
            .map(|p| p.join("data"))
            .unwrap_or_else(|| PathBuf::from("data"));

        if let Ok(params) = loader.load_from_spec_file(&data_dir) {
            for param in params {
                let _ = registry.register(param);
            }
        }

        SearchParameterExtractor::new(Arc::new(RwLock::new(registry)))
    }

    #[test]
    fn test_extract_patient_name() {
        let extractor = create_test_extractor();

        let patient = json!({
            "resourceType": "Patient",
            "id": "123",
            "name": [
                {
                    "family": "Smith",
                    "given": ["John", "James"]
                }
            ]
        });

        let values = extractor.extract(&patient, "Patient").unwrap();

        // Should have extracted name values
        let name_values: Vec<_> = values.iter().filter(|v| v.param_name == "name").collect();
        assert!(!name_values.is_empty(), "Should extract 'name' values");

        // Should have extracted family
        let family_values: Vec<_> = values.iter().filter(|v| v.param_name == "family").collect();
        assert!(!family_values.is_empty(), "Should extract 'family' values");
    }

    #[test]
    fn test_extract_patient_identifier() {
        let extractor = create_test_extractor();

        let patient = json!({
            "resourceType": "Patient",
            "id": "123",
            "identifier": [
                {
                    "system": "http://hospital.org/mrn",
                    "value": "12345"
                }
            ]
        });

        let values = extractor.extract(&patient, "Patient").unwrap();

        let id_values: Vec<_> = values
            .iter()
            .filter(|v| v.param_name == "identifier")
            .collect();
        assert!(!id_values.is_empty(), "Should extract 'identifier' values");

        if let IndexValue::Token { system, code, .. } = &id_values[0].value {
            assert_eq!(system.as_ref().unwrap(), "http://hospital.org/mrn");
            assert_eq!(code, "12345");
        }
    }

    #[test]
    fn test_extract_observation_values() {
        let extractor = create_test_extractor();

        let observation = json!({
            "resourceType": "Observation",
            "id": "obs1",
            "code": {
                "coding": [
                    {
                        "system": "http://loinc.org",
                        "code": "8867-4"
                    }
                ]
            },
            "subject": {
                "reference": "Patient/123"
            },
            "valueQuantity": {
                "value": 120.5,
                "unit": "mmHg"
            }
        });

        let values = extractor.extract(&observation, "Observation").unwrap();

        // Should have code
        let code_values: Vec<_> = values.iter().filter(|v| v.param_name == "code").collect();
        assert!(!code_values.is_empty(), "Should extract 'code' values");

        // Should have subject
        let subject_values: Vec<_> = values
            .iter()
            .filter(|v| v.param_name == "subject")
            .collect();
        assert!(
            !subject_values.is_empty(),
            "Should extract 'subject' values"
        );
    }

    #[test]
    fn test_invalid_resource() {
        let extractor = create_test_extractor();

        let not_object = json!("string");
        let result = extractor.extract(&not_object, "Patient");
        assert!(result.is_err());
    }

    #[test]
    fn test_resource_type_mismatch() {
        let extractor = create_test_extractor();

        let patient = json!({
            "resourceType": "Patient",
            "id": "123"
        });

        let result = extractor.extract(&patient, "Observation");
        assert!(result.is_err());
    }

    #[test]
    fn test_fhirpath_with_where_clause() {
        let extractor = create_test_extractor();

        // Test a patient with multiple names - FHIRPath should be able to filter
        let patient = json!({
            "resourceType": "Patient",
            "id": "123",
            "name": [
                {
                    "use": "official",
                    "family": "Smith",
                    "given": ["John"]
                },
                {
                    "use": "nickname",
                    "given": ["Johnny"]
                }
            ]
        });

        let values = extractor.extract(&patient, "Patient").unwrap();

        // Should extract all names (both official and nickname)
        let name_values: Vec<_> = values.iter().filter(|v| v.param_name == "name").collect();
        assert!(
            name_values.len() >= 2,
            "Should extract multiple name values"
        );
    }

    #[test]
    fn test_extract_observation_code_with_display() {
        let extractor = create_test_extractor();

        let observation = json!({
            "resourceType": "Observation",
            "id": "obs1",
            "status": "final",
            "code": {
                "coding": [
                    {
                        "system": "http://loinc.org",
                        "code": "8867-4",
                        "display": "Heart rate"
                    }
                ]
            }
        });

        // Extract values
        let values = extractor.extract(&observation, "Observation").unwrap();

        // Should have extracted code values
        let code_values: Vec<_> = values.iter().filter(|v| v.param_name == "code").collect();
        assert!(!code_values.is_empty(), "Should extract 'code' values");

        // Check that display is populated
        if let Some(first_code) = code_values.first() {
            if let IndexValue::Token { display, .. } = &first_code.value {
                assert_eq!(
                    display.as_deref(),
                    Some("Heart rate"),
                    "Display should be populated"
                );
            }
        }
    }

    #[test]
    fn test_extract_resource_id() {
        let extractor = create_test_extractor();

        let patient = json!({
            "resourceType": "Patient",
            "id": "p1"
        });

        let values = extractor.extract(&patient, "Patient").unwrap();

        // Should have extracted _id
        let id_values: Vec<_> = values.iter().filter(|v| v.param_name == "_id").collect();
        assert!(!id_values.is_empty(), "Should extract '_id' parameter");

        // Check the value
        if let Some(first_id) = id_values.first() {
            if let IndexValue::Token { code, .. } = &first_id.value {
                assert_eq!(code, "p1", "_id should be 'p1'");
            }
        }
    }

    #[test]
    fn test_json_to_evaluation_result() {
        // Test basic types
        assert!(matches!(
            json_to_evaluation_result(&json!(null)).unwrap(),
            EvaluationResult::Empty
        ));

        assert!(matches!(
            json_to_evaluation_result(&json!(true)).unwrap(),
            EvaluationResult::Boolean(true, _)
        ));

        assert!(matches!(
            json_to_evaluation_result(&json!("test")).unwrap(),
            EvaluationResult::String(s, _) if s == "test"
        ));

        assert!(matches!(
            json_to_evaluation_result(&json!(42)).unwrap(),
            EvaluationResult::Integer(42, _)
        ));

        // Test array
        if let EvaluationResult::Collection { items, .. } =
            json_to_evaluation_result(&json!([1, 2, 3])).unwrap()
        {
            assert_eq!(items.len(), 3);
        } else {
            panic!("Expected collection");
        }

        // Test object
        if let EvaluationResult::Object { map, .. } =
            json_to_evaluation_result(&json!({"key": "value"})).unwrap()
        {
            assert!(map.contains_key("key"));
        } else {
            panic!("Expected object");
        }
    }

    #[test]
    fn test_filter_expression_for_resource() {
        let extractor = create_test_extractor();

        // Test multi-resource expression (like patient search param)
        let complex_expr =
            "AllergyIntolerance.patient | Immunization.patient | Observation.subject";
        let filtered = extractor.filter_expression_for_resource(complex_expr, "Immunization");
        assert_eq!(filtered, "Immunization.patient");

        // Test with no matching parts - should return original
        let no_match = extractor.filter_expression_for_resource(complex_expr, "Patient");
        assert_eq!(no_match, complex_expr);

        // Test simple expression (single resource type)
        let simple_expr = "Patient.name";
        let simple_filtered = extractor.filter_expression_for_resource(simple_expr, "Patient");
        assert_eq!(simple_filtered, "Patient.name");

        // Test that partial matches don't count (Observation shouldn't match Obs)
        let partial = extractor.filter_expression_for_resource("Observation.code", "Obs");
        assert_eq!(partial, "Observation.code");

        // Test stripping .where(resolve() is X) pattern
        let with_resolve = "Observation.subject.where(resolve() is Patient) | Patient.link.other";
        let stripped = extractor.filter_expression_for_resource(with_resolve, "Observation");
        assert_eq!(stripped, "Observation.subject");

        // Test real-world patient search param pattern
        let patient_expr = "CarePlan.subject.where(resolve() is Patient) | Observation.subject.where(resolve() is Patient)";
        let careplan_filtered = extractor.filter_expression_for_resource(patient_expr, "CarePlan");
        assert_eq!(careplan_filtered, "CarePlan.subject");
        let obs_filtered = extractor.filter_expression_for_resource(patient_expr, "Observation");
        assert_eq!(obs_filtered, "Observation.subject");
    }

    #[test]
    fn test_extract_immunization_patient() {
        let extractor = create_test_extractor();

        let immunization = json!({
            "resourceType": "Immunization",
            "id": "test-imm",
            "status": "completed",
            "vaccineCode": {
                "coding": [{
                    "system": "http://hl7.org/fhir/sid/cvx",
                    "code": "140"
                }]
            },
            "patient": {
                "reference": "Patient/test-patient"
            },
            "occurrenceDateTime": "2021-01-01"
        });

        let values = extractor.extract(&immunization, "Immunization").unwrap();

        // Should have extracted patient reference
        let patient_values: Vec<_> = values
            .iter()
            .filter(|v| v.param_name == "patient")
            .collect();
        assert!(
            !patient_values.is_empty(),
            "Should extract 'patient' values from Immunization"
        );

        // Check the reference value
        if let IndexValue::Reference { reference, .. } = &patient_values[0].value {
            assert!(
                reference.contains("Patient/test-patient") || reference.contains("test-patient"),
                "Should contain patient reference, got: {}",
                reference
            );
        }
    }
}