ahri_tre_tabular/

metadata.rs

use std::collections::{BTreeMap, BTreeSet, HashMap};

use arrow_schema::{Field, Schema};
use serde_json::Value;

use crate::{ArrowTable, Result};

pub const AHRI_TRE_METADATA_PREFIX: &str = "ahri_tre.";

pub mod keys {
    pub const DATASET_ID: &str = "ahri_tre.dataset_id";
    pub const ASSET_ID: &str = "ahri_tre.asset_id";
    pub const ASSET_NAME: &str = "ahri_tre.asset_name";
    pub const STUDY_ID: &str = "ahri_tre.study_id";
    pub const VERSION_ID: &str = "ahri_tre.version_id";
    pub const VERSION: &str = "ahri_tre.version";
    pub const DESCRIPTION: &str = "ahri_tre.description";
    pub const VERSION_NOTE: &str = "ahri_tre.version_note";

    pub const VARIABLE_ID: &str = "ahri_tre.variable_id";
    pub const VARIABLE_NAME: &str = "ahri_tre.variable_name";
    pub const VALUE_TYPE_ID: &str = "ahri_tre.value_type_id";
    pub const VALUE_TYPE: &str = "ahri_tre.value_type";
    pub const VALUE_FORMAT: &str = "ahri_tre.value_format";
    pub const ROW_ROLE: &str = "ahri_tre.row_role";
    pub const NOTE: &str = "ahri_tre.note";
    pub const ONTOLOGY_NAMESPACE: &str = "ahri_tre.ontology_namespace";
    pub const ONTOLOGY_CLASS: &str = "ahri_tre.ontology_class";
    pub const VOCABULARY_ID: &str = "ahri_tre.vocabulary_id";
    pub const VOCABULARY_NAME: &str = "ahri_tre.vocabulary_name";
    pub const VOCABULARY_DESCRIPTION: &str = "ahri_tre.vocabulary_description";
    pub const VOCABULARY_ITEMS: &str = "ahri_tre.vocabulary_items";
}

const CONVENTIONAL_DESCRIPTION_KEYS: [&str; 4] = ["description", "title", "label", "comment"];
const CONVENTIONAL_VOCABULARY_KEYS: [&str; 6] = [
    "value_labels",
    "value_descriptions",
    "code_labels",
    "code_descriptions",
    "codes",
    "categories",
];

#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct AhriTreArrowMetadata {
    pub schema: BTreeMap<String, String>,
    pub fields: Vec<AhriTreFieldMetadata>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AhriTreFieldMetadata {
    pub field_name: String,
    pub metadata: BTreeMap<String, String>,
}

#[derive(Debug, Clone)]
pub struct AnnotatedArrowTable {
    pub table: ArrowTable,
    pub warnings: Vec<String>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FieldMetadataProposal {
    pub field_name: String,
    pub variable_name: Option<String>,
    pub value_type_id: Option<i64>,
    pub value_type: Option<String>,
    pub description: Option<String>,
    pub vocabulary_name: Option<String>,
    pub vocabulary_description: Option<String>,
    pub vocabulary_items: Vec<VocabularyItemProposal>,
    pub warnings: Vec<String>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct VocabularyItemProposal {
    pub code: String,
    pub description: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct JsonMetadataExtraction {
    pub proposals: Vec<FieldMetadataProposal>,
    pub warnings: Vec<String>,
}

pub fn annotate_arrow_table(
    table: &ArrowTable,
    metadata: &AhriTreArrowMetadata,
) -> Result<AnnotatedArrowTable> {
    let (schema, warnings) = annotate_arrow_schema(table.schema().as_ref(), metadata);
    let table = table.with_schema(schema)?;
    Ok(AnnotatedArrowTable { table, warnings })
}

pub fn annotate_arrow_schema(
    schema: &Schema,
    metadata: &AhriTreArrowMetadata,
) -> (Schema, Vec<String>) {
    let mut schema_metadata = without_ahri_tre_metadata(schema.metadata());
    for (key, value) in &metadata.schema {
        schema_metadata.insert(key.clone(), value.clone());
    }

    let field_metadata_index = FieldMetadataIndex::new(&metadata.fields);

    let mut warnings = Vec::new();
    let mut matched_fields = BTreeSet::new();
    let fields = schema
        .fields()
        .iter()
        .map(|field| {
            let field = field.as_ref();
            let mut merged = without_ahri_tre_metadata(field.metadata());
            if let Some((field_index, tre_metadata)) =
                field_metadata_index.find(field.name(), &matched_fields, &mut warnings)
            {
                for (key, value) in tre_metadata.iter() {
                    merged.insert(key.clone(), value.clone());
                }
                matched_fields.insert(field_index);
            } else {
                warnings.push(format!(
                    "Arrow field {} has no registered AHRI_TRE variable metadata",
                    field.name()
                ));
            }
            clone_field_with_metadata(field, merged)
        })
        .collect::<Vec<_>>();

    for (field_index, field_metadata) in metadata.fields.iter().enumerate() {
        if !matched_fields.contains(&field_index) {
            warnings.push(format!(
                "registered AHRI_TRE variable {} is not present in the Arrow schema",
                field_metadata.field_name
            ));
        }
    }

    (Schema::new_with_metadata(fields, schema_metadata), warnings)
}

struct FieldMetadataIndex<'a> {
    fields: &'a [AhriTreFieldMetadata],
    exact: BTreeMap<&'a str, Vec<usize>>,
    canonical: BTreeMap<String, Vec<usize>>,
}

impl<'a> FieldMetadataIndex<'a> {
    fn new(fields: &'a [AhriTreFieldMetadata]) -> Self {
        let mut exact = BTreeMap::new();
        let mut canonical = BTreeMap::new();
        for (index, field_metadata) in fields.iter().enumerate() {
            exact
                .entry(field_metadata.field_name.as_str())
                .or_insert_with(Vec::new)
                .push(index);
            canonical
                .entry(canonical_field_name(&field_metadata.field_name))
                .or_insert_with(Vec::new)
                .push(index);
        }
        Self {
            fields,
            exact,
            canonical,
        }
    }

    fn find(
        &self,
        arrow_field_name: &str,
        matched_fields: &BTreeSet<usize>,
        warnings: &mut Vec<String>,
    ) -> Option<(usize, &'a BTreeMap<String, String>)> {
        if let Some(indexes) = self.exact.get(arrow_field_name) {
            return self.unique_match(
                arrow_field_name,
                indexes,
                matched_fields,
                "exact field name",
                warnings,
            );
        }

        let canonical_name = canonical_field_name(arrow_field_name);
        let indexes = self.canonical.get(&canonical_name)?;
        self.unique_match(
            arrow_field_name,
            indexes,
            matched_fields,
            "canonical field name",
            warnings,
        )
    }

    fn unique_match(
        &self,
        arrow_field_name: &str,
        indexes: &[usize],
        matched_fields: &BTreeSet<usize>,
        match_kind: &str,
        warnings: &mut Vec<String>,
    ) -> Option<(usize, &'a BTreeMap<String, String>)> {
        let unmatched = indexes
            .iter()
            .copied()
            .filter(|index| !matched_fields.contains(index))
            .collect::<Vec<_>>();
        match unmatched.as_slice() {
            [index] => Some((*index, &self.fields[*index].metadata)),
            [] => None,
            _ => {
                warnings.push(format!(
                    "Arrow field {arrow_field_name} has ambiguous registered AHRI_TRE variable metadata by {match_kind}"
                ));
                None
            }
        }
    }
}

fn canonical_field_name(field_name: &str) -> String {
    field_name.to_ascii_lowercase()
}

pub fn extract_field_metadata_proposals(schema: &Schema) -> Vec<FieldMetadataProposal> {
    schema
        .fields()
        .iter()
        .map(|field| extract_field_metadata_proposal(field.as_ref()))
        .collect()
}

pub fn extract_field_metadata_proposal(field: &Field) -> FieldMetadataProposal {
    let metadata = field.metadata();
    let description = metadata
        .get(keys::DESCRIPTION)
        .cloned()
        .or_else(|| conventional_description(metadata));
    let mut warnings = Vec::new();
    let vocabulary_items = extract_vocabulary_items(metadata).unwrap_or_else(|metadata_warnings| {
        warnings.extend(metadata_warnings);
        Vec::new()
    });

    FieldMetadataProposal {
        field_name: field.name().clone(),
        variable_name: metadata.get(keys::VARIABLE_NAME).cloned(),
        value_type_id: metadata
            .get(keys::VALUE_TYPE_ID)
            .and_then(|value| value.trim().parse::<i64>().ok()),
        value_type: metadata.get(keys::VALUE_TYPE).cloned(),
        description,
        vocabulary_name: metadata.get(keys::VOCABULARY_NAME).cloned(),
        vocabulary_description: metadata.get(keys::VOCABULARY_DESCRIPTION).cloned(),
        vocabulary_items,
        warnings,
    }
}

pub fn extract_json_field_metadata_proposals(text: &str) -> JsonMetadataExtraction {
    let Ok(root) = serde_json::from_str::<Value>(text) else {
        return JsonMetadataExtraction::default();
    };
    let Value::Object(root) = root else {
        return JsonMetadataExtraction::default();
    };

    let common = extract_common_json_metadata(&root);
    let ahri_tre = extract_ahri_tre_json_metadata(&root);
    merge_json_metadata(common, ahri_tre)
}

fn extract_ahri_tre_json_metadata(root: &serde_json::Map<String, Value>) -> JsonMetadataExtraction {
    let Some(contract) = root
        .get("ahri_tre")
        .or_else(|| root.get("ahri_tre_metadata"))
    else {
        return JsonMetadataExtraction::default();
    };

    let Some(fields) = contract
        .get("fields")
        .or_else(|| contract.get("variables"))
        .and_then(Value::as_array)
    else {
        return JsonMetadataExtraction {
            proposals: Vec::new(),
            warnings: vec![
                "ignored AHRI_TRE JSON metadata: expected ahri_tre.fields array".to_string(),
            ],
        };
    };

    let mut extraction = JsonMetadataExtraction::default();
    for field in fields {
        let Value::Object(field) = field else {
            extraction
                .warnings
                .push("ignored AHRI_TRE JSON metadata field: expected object".to_string());
            continue;
        };
        let Some(field_name) = field
            .get("field_name")
            .or_else(|| field.get("name"))
            .and_then(Value::as_str)
            .filter(|value| !value.trim().is_empty())
        else {
            extraction
                .warnings
                .push("ignored AHRI_TRE JSON metadata field: missing field_name".to_string());
            continue;
        };
        let mut proposal = empty_field_metadata_proposal(field_name);
        proposal.variable_name = field
            .get("variable_name")
            .or_else(|| field.get("tre_variable_name"))
            .and_then(json_value_string_ref);
        proposal.value_type_id = field
            .get("value_type_id")
            .and_then(Value::as_i64)
            .or_else(|| {
                field
                    .get("value_type_id")
                    .and_then(Value::as_str)
                    .and_then(|value| value.trim().parse::<i64>().ok())
            });
        proposal.value_type = field
            .get("value_type")
            .or_else(|| field.get("type"))
            .and_then(json_value_string_ref);
        proposal.description = field
            .get("description")
            .or_else(|| field.get("label"))
            .and_then(json_value_string_ref);
        proposal.vocabulary_name = field.get("vocabulary_name").and_then(json_value_string_ref);
        proposal.vocabulary_description = field
            .get("vocabulary_description")
            .and_then(json_value_string_ref);
        if let Some(value) = field
            .get("vocabulary_items")
            .or_else(|| field.get("categories"))
            .or_else(|| field.get("codes"))
        {
            match parse_vocabulary_items_value("AHRI_TRE JSON vocabulary_items", value.clone()) {
                Ok(items) => proposal.vocabulary_items = items,
                Err(warnings) => proposal.warnings.extend(warnings),
            }
        }
        extraction.proposals.push(proposal);
    }
    extraction
}

fn extract_common_json_metadata(root: &serde_json::Map<String, Value>) -> JsonMetadataExtraction {
    if let Some(extraction) = extract_json_schema_metadata(root)
        && (!extraction.proposals.is_empty() || !extraction.warnings.is_empty())
    {
        return extraction;
    }
    JsonMetadataExtraction::default()
}

fn extract_json_schema_metadata(
    root: &serde_json::Map<String, Value>,
) -> Option<JsonMetadataExtraction> {
    let properties = root.get("properties")?.as_object()?;
    let mut extraction = JsonMetadataExtraction::default();
    for (field_name, descriptor) in properties {
        let Value::Object(descriptor) = descriptor else {
            extraction.warnings.push(format!(
                "ignored JSON Schema property {field_name}: expected object descriptor"
            ));
            continue;
        };
        let mut proposal = empty_field_metadata_proposal(field_name);
        proposal.value_type = descriptor
            .get("type")
            .and_then(json_schema_type_name)
            .or_else(|| {
                descriptor
                    .get("format")
                    .and_then(Value::as_str)
                    .and_then(json_schema_format_type_name)
            });
        proposal.description = descriptor
            .get("description")
            .or_else(|| descriptor.get("title"))
            .and_then(json_value_string_ref);

        if let Some(items) = json_schema_one_of_vocabulary_items(descriptor) {
            proposal.vocabulary_items = items;
            proposal.value_type = Some("enumeration".to_string());
        } else if let Some(items) = json_schema_enum_names_vocabulary_items(descriptor) {
            proposal.vocabulary_items = items;
            proposal.value_type = Some("enumeration".to_string());
        } else if descriptor.get("enum").is_some() {
            proposal.warnings.push(format!(
                "ignored JSON Schema property {field_name} enum: expected enumNames or oneOf const labels for categorical metadata"
            ));
        }
        extraction.proposals.push(proposal);
    }
    Some(extraction)
}

fn merge_json_metadata(
    common: JsonMetadataExtraction,
    ahri_tre: JsonMetadataExtraction,
) -> JsonMetadataExtraction {
    let mut proposals_by_field = BTreeMap::new();
    for proposal in common.proposals {
        proposals_by_field.insert(proposal.field_name.clone(), proposal);
    }
    for ahri_proposal in ahri_tre.proposals {
        proposals_by_field
            .entry(ahri_proposal.field_name.clone())
            .and_modify(|common_proposal| {
                overlay_ahri_tre_proposal(common_proposal, &ahri_proposal)
            })
            .or_insert(ahri_proposal);
    }
    let mut warnings = common.warnings;
    warnings.extend(ahri_tre.warnings);
    JsonMetadataExtraction {
        proposals: proposals_by_field.into_values().collect(),
        warnings,
    }
}

fn overlay_ahri_tre_proposal(common: &mut FieldMetadataProposal, ahri_tre: &FieldMetadataProposal) {
    if ahri_tre.variable_name.is_some() {
        common.variable_name = ahri_tre.variable_name.clone();
    }
    if ahri_tre.value_type_id.is_some() {
        common.value_type_id = ahri_tre.value_type_id;
    }
    if ahri_tre.value_type.is_some() {
        common.value_type = ahri_tre.value_type.clone();
    }
    if ahri_tre.description.is_some() {
        common.description = ahri_tre.description.clone();
    }
    if ahri_tre.vocabulary_name.is_some() {
        common.vocabulary_name = ahri_tre.vocabulary_name.clone();
    }
    if ahri_tre.vocabulary_description.is_some() {
        common.vocabulary_description = ahri_tre.vocabulary_description.clone();
    }
    if !ahri_tre.vocabulary_items.is_empty() {
        common.vocabulary_items = ahri_tre.vocabulary_items.clone();
    }
    common.warnings.extend(ahri_tre.warnings.clone());
}

fn empty_field_metadata_proposal(field_name: &str) -> FieldMetadataProposal {
    FieldMetadataProposal {
        field_name: field_name.to_string(),
        variable_name: None,
        value_type_id: None,
        value_type: None,
        description: None,
        vocabulary_name: None,
        vocabulary_description: None,
        vocabulary_items: Vec::new(),
        warnings: Vec::new(),
    }
}

fn without_ahri_tre_metadata(metadata: &HashMap<String, String>) -> HashMap<String, String> {
    metadata
        .iter()
        .filter(|(key, _)| !key.starts_with(AHRI_TRE_METADATA_PREFIX))
        .map(|(key, value)| (key.clone(), value.clone()))
        .collect()
}

fn clone_field_with_metadata(field: &Field, metadata: HashMap<String, String>) -> Field {
    field.clone().with_metadata(metadata)
}

fn conventional_description(metadata: &HashMap<String, String>) -> Option<String> {
    CONVENTIONAL_DESCRIPTION_KEYS
        .iter()
        .find_map(|key| metadata.get(*key).cloned())
}

fn extract_vocabulary_items(
    metadata: &HashMap<String, String>,
) -> std::result::Result<Vec<VocabularyItemProposal>, Vec<String>> {
    if let Some(value) = metadata.get(keys::VOCABULARY_ITEMS) {
        return parse_vocabulary_items(keys::VOCABULARY_ITEMS, value);
    }

    for key in CONVENTIONAL_VOCABULARY_KEYS {
        if let Some(value) = metadata.get(key) {
            return parse_vocabulary_items(key, value);
        }
    }
    Ok(Vec::new())
}

fn parse_vocabulary_items(
    key: &str,
    value: &str,
) -> std::result::Result<Vec<VocabularyItemProposal>, Vec<String>> {
    let parsed = serde_json::from_str::<Value>(value).map_err(|error| {
        vec![format!(
            "ignored metadata key {key}: expected JSON code mapping ({error})"
        )]
    })?;
    parse_vocabulary_items_value(key, parsed)
}

fn parse_vocabulary_items_value(
    key: &str,
    parsed: Value,
) -> std::result::Result<Vec<VocabularyItemProposal>, Vec<String>> {
    let items = match parsed {
        Value::Object(map) => map
            .into_iter()
            .filter_map(|(code, value)| {
                description_from_json_value(value).map(|description| VocabularyItemProposal {
                    code,
                    description: Some(description),
                })
            })
            .collect::<Vec<_>>(),
        Value::Array(values) => {
            let mut items = Vec::with_capacity(values.len());
            for value in values {
                let Value::Object(mut object) = value else {
                    return Err(vec![format!(
                        "ignored metadata key {key}: array items must be objects with code and label or description"
                    )]);
                };
                let Some(code) = object.remove("code").and_then(json_string) else {
                    return Err(vec![format!(
                        "ignored metadata key {key}: vocabulary item is missing code"
                    )]);
                };
                let description = object
                    .remove("label")
                    .or_else(|| object.remove("description"))
                    .and_then(description_from_json_value);
                items.push(VocabularyItemProposal { code, description });
            }
            items
        }
        _ => {
            return Err(vec![format!(
                "ignored metadata key {key}: expected JSON object or array"
            )]);
        }
    };
    if items.is_empty() {
        Err(vec![format!(
            "ignored metadata key {key}: no vocabulary items found"
        )])
    } else {
        Ok(items)
    }
}

fn json_schema_type_name(value: &Value) -> Option<String> {
    match value {
        Value::String(value) => match value.as_str() {
            "integer" => Some("integer".to_string()),
            "number" => Some("decimal".to_string()),
            "string" => Some("string".to_string()),
            "boolean" => Some("boolean".to_string()),
            _ => None,
        },
        Value::Array(values) => values.iter().find_map(json_schema_type_name),
        _ => None,
    }
}

fn json_schema_format_type_name(value: &str) -> Option<String> {
    match value {
        "date" => Some("date".to_string()),
        "date-time" => Some("datetime".to_string()),
        "time" => Some("time".to_string()),
        _ => None,
    }
}

fn json_schema_enum_names_vocabulary_items(
    descriptor: &serde_json::Map<String, Value>,
) -> Option<Vec<VocabularyItemProposal>> {
    let enums = descriptor.get("enum")?.as_array()?;
    let names = descriptor
        .get("enumNames")
        .or_else(|| descriptor.get("enumDescriptions"))?
        .as_array()?;
    if enums.len() != names.len() || enums.is_empty() {
        return None;
    }
    let mut items = Vec::with_capacity(enums.len());
    for (code, description) in enums.iter().zip(names) {
        items.push(VocabularyItemProposal {
            code: json_value_string_ref(code)?,
            description: Some(json_value_string_ref(description)?),
        });
    }
    Some(items)
}

fn json_schema_one_of_vocabulary_items(
    descriptor: &serde_json::Map<String, Value>,
) -> Option<Vec<VocabularyItemProposal>> {
    let one_of = descriptor
        .get("oneOf")
        .or_else(|| descriptor.get("anyOf"))?
        .as_array()?;
    if one_of.is_empty() {
        return None;
    }
    let mut items = Vec::with_capacity(one_of.len());
    for item in one_of {
        let item = item.as_object()?;
        let code = item
            .get("const")
            .or_else(|| item.get("enum").and_then(|value| value.as_array()?.first()))
            .and_then(json_value_string_ref)?;
        let description = item
            .get("title")
            .or_else(|| item.get("description"))
            .and_then(json_value_string_ref)?;
        items.push(VocabularyItemProposal {
            code,
            description: Some(description),
        });
    }
    Some(items)
}

fn description_from_json_value(value: Value) -> Option<String> {
    match value {
        Value::String(value) => Some(value),
        Value::Object(mut object) => object
            .remove("label")
            .or_else(|| object.remove("description"))
            .and_then(json_string),
        _ => None,
    }
}

fn json_string(value: Value) -> Option<String> {
    match value {
        Value::String(value) => Some(value),
        Value::Number(value) => Some(value.to_string()),
        Value::Bool(value) => Some(value.to_string()),
        _ => None,
    }
}

fn json_value_string_ref(value: &Value) -> Option<String> {
    match value {
        Value::String(value) => Some(value.clone()),
        Value::Number(value) => Some(value.to_string()),
        Value::Bool(value) => Some(value.to_string()),
        _ => None,
    }
}

/// Scans large JSON tables without retaining rows. Only the existing metadata
/// keys are retained, each under the fixed metadata allocation input ceiling.
pub fn extract_json_file_metadata_proposals(
    path: &std::path::Path,
) -> std::io::Result<JsonMetadataExtraction> {
    use serde::Deserializer;
    use serde::de::{MapAccess, Visitor};
    use std::{
        cell::Cell,
        io::{BufReader, Read},
        rc::Rc,
    };
    const LIMIT: usize = crate::streaming::MAX_BATCH_BYTES;
    struct BudgetReader<R> {
        inner: R,
        remaining: Rc<Cell<Option<usize>>>,
    }
    impl<R: Read> Read for BudgetReader<R> {
        fn read(&mut self, out: &mut [u8]) -> std::io::Result<usize> {
            if let Some(left) = self.remaining.get() {
                if left == 0 {
                    return Err(std::io::Error::other("JSON metadata exceeds its limit"));
                }
                let len = out.len().min(left);
                let count = self.inner.read(&mut out[..len])?;
                self.remaining.set(Some(left - count));
                Ok(count)
            } else {
                self.inner.read(out)
            }
        }
    }
    struct Metadata(Rc<Cell<Option<usize>>>);
    impl<'de> Visitor<'de> for Metadata {
        type Value = serde_json::Map<String, Value>;
        fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            f.write_str("JSON object")
        }
        fn visit_map<M: MapAccess<'de>>(
            self,
            mut map: M,
        ) -> std::result::Result<Self::Value, M::Error> {
            let mut metadata = serde_json::Map::new();
            let mut retained = 0usize;
            loop {
                self.0.set(Some(LIMIT - retained));
                let Some(key) = map.next_key::<String>()? else {
                    break;
                };
                if matches!(
                    key.as_str(),
                    "properties" | "ahri_tre" | "ahri_tre_metadata"
                ) {
                    let value = map.next_value::<Value>()?;
                    retained = LIMIT - self.0.get().unwrap_or(0);
                    metadata.insert(key, value);
                } else {
                    self.0.set(None);
                    map.next_value::<serde::de::IgnoredAny>()?;
                }
            }
            Ok(metadata)
        }
    }
    let remaining = Rc::new(Cell::new(Some(LIMIT)));
    let reader = BudgetReader {
        inner: BufReader::new(std::fs::File::open(path)?),
        remaining: remaining.clone(),
    };
    let mut deserializer = serde_json::Deserializer::from_reader(reader);
    let root = match deserializer.deserialize_map(Metadata(remaining)) {
        Ok(root) => root,
        Err(error) if error.is_io() => {
            return Err(std::io::Error::other("JSON metadata exceeds its limit"));
        }
        Err(_) => return Ok(JsonMetadataExtraction::default()),
    };
    Ok(merge_json_metadata(
        extract_common_json_metadata(&root),
        extract_ahri_tre_json_metadata(&root),
    ))
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow_schema::DataType;

    #[test]
    fn ahri_tre_description_wins_over_conventional_description() {
        let field = Field::new("sex", DataType::Utf8, true).with_metadata(
            [
                (
                    "description".to_string(),
                    "Conventional description".to_string(),
                ),
                (
                    keys::DESCRIPTION.to_string(),
                    "AHRI TRE description".to_string(),
                ),
            ]
            .into_iter()
            .collect(),
        );

        let proposal = extract_field_metadata_proposal(&field);

        assert_eq!(
            proposal.description.as_deref(),
            Some("AHRI TRE description")
        );
    }

    #[test]
    fn annotation_matches_registered_metadata_by_canonical_field_name() {
        let schema = Schema::new(vec![Field::new("HouseholdId", DataType::Utf8, true)]);
        let metadata = AhriTreArrowMetadata {
            schema: BTreeMap::new(),
            fields: vec![AhriTreFieldMetadata {
                field_name: "householdid".to_string(),
                metadata: [
                    (keys::VARIABLE_NAME.to_string(), "householdid".to_string()),
                    (keys::VALUE_TYPE_ID.to_string(), "7".to_string()),
                ]
                .into_iter()
                .collect(),
            }],
        };

        let (annotated_schema, warnings) = annotate_arrow_schema(&schema, &metadata);

        assert!(warnings.is_empty(), "{warnings:?}");
        let field = annotated_schema.field_with_name("HouseholdId").unwrap();
        assert_eq!(
            field.metadata()[keys::VARIABLE_NAME],
            "householdid",
            "canonical variable metadata should attach to the original Arrow field name"
        );
    }

    #[test]
    fn annotation_warns_instead_of_guessing_ambiguous_canonical_metadata() {
        let schema = Schema::new(vec![Field::new("HouseholdId", DataType::Utf8, true)]);
        let metadata = AhriTreArrowMetadata {
            schema: BTreeMap::new(),
            fields: vec![
                AhriTreFieldMetadata {
                    field_name: "householdid".to_string(),
                    metadata: [(keys::VARIABLE_NAME.to_string(), "householdid".to_string())]
                        .into_iter()
                        .collect(),
                },
                AhriTreFieldMetadata {
                    field_name: "HOUSEHOLDID".to_string(),
                    metadata: [(keys::VARIABLE_NAME.to_string(), "HOUSEHOLDID".to_string())]
                        .into_iter()
                        .collect(),
                },
            ],
        };

        let (annotated_schema, warnings) = annotate_arrow_schema(&schema, &metadata);

        assert!(
            warnings
                .iter()
                .any(|warning| warning.contains("ambiguous registered AHRI_TRE variable metadata")),
            "{warnings:?}"
        );
        assert!(
            !annotated_schema
                .field_with_name("HouseholdId")
                .unwrap()
                .metadata()
                .contains_key(keys::VARIABLE_NAME)
        );
    }

    #[test]
    fn conventional_description_keys_are_recognized_but_unrelated_metadata_is_ignored() {
        let field = Field::new("age", DataType::Int64, true).with_metadata(
            [
                ("title".to_string(), "Age at visit".to_string()),
                ("owner".to_string(), "not TRE semantics".to_string()),
            ]
            .into_iter()
            .collect(),
        );

        let proposal = extract_field_metadata_proposal(&field);

        assert_eq!(proposal.description.as_deref(), Some("Age at visit"));
        assert!(proposal.variable_name.is_none());
        assert!(proposal.vocabulary_items.is_empty());
    }

    #[test]
    fn ahri_tre_vocabulary_items_preserve_codes_and_order() {
        let field = Field::new("status", DataType::Utf8, true).with_metadata(
            [(
                keys::VOCABULARY_ITEMS.to_string(),
                r#"[{"code":"A","label":"Alive"},{"code":"D","description":"Dead"}]"#.to_string(),
            )]
            .into_iter()
            .collect(),
        );

        let proposal = extract_field_metadata_proposal(&field);

        assert_eq!(
            proposal.vocabulary_items,
            vec![
                VocabularyItemProposal {
                    code: "A".to_string(),
                    description: Some("Alive".to_string()),
                },
                VocabularyItemProposal {
                    code: "D".to_string(),
                    description: Some("Dead".to_string()),
                },
            ]
        );
        assert!(proposal.warnings.is_empty());
    }

    #[test]
    fn conventional_code_label_mapping_creates_vocabulary_items() {
        let field = Field::new("sex", DataType::Utf8, true).with_metadata(
            [(
                "value_labels".to_string(),
                r#"{"F":"Female","M":"Male"}"#.to_string(),
            )]
            .into_iter()
            .collect(),
        );

        let proposal = extract_field_metadata_proposal(&field);

        assert_eq!(
            proposal.vocabulary_items,
            vec![
                VocabularyItemProposal {
                    code: "F".to_string(),
                    description: Some("Female".to_string()),
                },
                VocabularyItemProposal {
                    code: "M".to_string(),
                    description: Some("Male".to_string()),
                },
            ]
        );
    }

    #[test]
    fn ahri_tre_json_metadata_extracts_field_contracts() {
        let extraction = extract_json_field_metadata_proposals(
            r#"{
                "ahri_tre": {
                    "fields": [{
                        "field_name": "status",
                        "variable_name": "status_code",
                        "value_type": "enumeration",
                        "description": "AHRI TRE status",
                        "vocabulary_name": "status_vocab",
                        "vocabulary_description": "Status code list",
                        "vocabulary_items": [
                            {"code": "A", "label": "Alive"},
                            {"code": "D", "description": "Dead"}
                        ]
                    }]
                }
            }"#,
        );

        assert!(extraction.warnings.is_empty());
        assert_eq!(extraction.proposals.len(), 1);
        let proposal = &extraction.proposals[0];
        assert_eq!(proposal.field_name, "status");
        assert_eq!(proposal.variable_name.as_deref(), Some("status_code"));
        assert_eq!(proposal.value_type.as_deref(), Some("enumeration"));
        assert_eq!(proposal.description.as_deref(), Some("AHRI TRE status"));
        assert_eq!(
            proposal.vocabulary_items,
            vec![
                VocabularyItemProposal {
                    code: "A".to_string(),
                    description: Some("Alive".to_string()),
                },
                VocabularyItemProposal {
                    code: "D".to_string(),
                    description: Some("Dead".to_string()),
                },
            ]
        );
    }

    #[test]
    fn json_schema_metadata_extracts_descriptions_types_and_one_of_vocabularies() {
        let extraction = extract_json_field_metadata_proposals(
            r#"{
                "$schema": "https://json-schema.org/draft/2020-12/schema",
                "type": "object",
                "properties": {
                    "age": {"type": "integer", "description": "Age in years"},
                    "sex": {
                        "type": "string",
                        "title": "Sex at baseline",
                        "oneOf": [
                            {"const": "F", "title": "Female"},
                            {"const": "M", "description": "Male"}
                        ]
                    }
                }
            }"#,
        );

        assert!(extraction.warnings.is_empty());
        let age = extraction
            .proposals
            .iter()
            .find(|proposal| proposal.field_name == "age")
            .expect("age proposal should exist");
        assert_eq!(age.value_type.as_deref(), Some("integer"));
        assert_eq!(age.description.as_deref(), Some("Age in years"));
        let sex = extraction
            .proposals
            .iter()
            .find(|proposal| proposal.field_name == "sex")
            .expect("sex proposal should exist");
        assert_eq!(sex.value_type.as_deref(), Some("enumeration"));
        assert_eq!(sex.description.as_deref(), Some("Sex at baseline"));
        assert_eq!(
            sex.vocabulary_items,
            vec![
                VocabularyItemProposal {
                    code: "F".to_string(),
                    description: Some("Female".to_string()),
                },
                VocabularyItemProposal {
                    code: "M".to_string(),
                    description: Some("Male".to_string()),
                },
            ]
        );
    }

    #[test]
    fn ahri_tre_json_metadata_overrides_common_schema_for_same_field() {
        let extraction = extract_json_field_metadata_proposals(
            r#"{
                "properties": {
                    "status": {
                        "type": "string",
                        "description": "Common status",
                        "enum": ["A", "D"],
                        "enumNames": ["Active", "Dormant"]
                    }
                },
                "ahri_tre": {
                    "fields": [{
                        "field_name": "status",
                        "variable_name": "status_code",
                        "description": "AHRI TRE status",
                        "vocabulary_items": {"A": "Alive", "D": "Dead"}
                    }]
                }
            }"#,
        );

        assert_eq!(extraction.proposals.len(), 1);
        let proposal = &extraction.proposals[0];
        assert_eq!(proposal.variable_name.as_deref(), Some("status_code"));
        assert_eq!(proposal.description.as_deref(), Some("AHRI TRE status"));
        assert_eq!(
            proposal
                .vocabulary_items
                .iter()
                .map(|item| (item.code.as_str(), item.description.as_deref()))
                .collect::<Vec<_>>(),
            vec![("A", Some("Alive")), ("D", Some("Dead"))]
        );
    }

    #[test]
    fn ordinary_json_data_is_not_interpreted_as_metadata() {
        let extraction = extract_json_field_metadata_proposals(
            r#"[{"status":"A","age":34},{"status":"D","age":41}]"#,
        );

        assert!(extraction.proposals.is_empty());
        assert!(extraction.warnings.is_empty());
    }

    #[test]
    fn ambiguous_json_schema_enum_is_reported_without_vocab_semantics() {
        let extraction = extract_json_field_metadata_proposals(
            r#"{
                "properties": {
                    "status": {
                        "type": "string",
                        "description": "Status",
                        "enum": ["A", "D"]
                    }
                }
            }"#,
        );

        assert_eq!(extraction.proposals.len(), 1);
        let proposal = &extraction.proposals[0];
        assert!(proposal.vocabulary_items.is_empty());
        assert!(
            proposal
                .warnings
                .iter()
                .any(|warning| warning.contains("expected enumNames or oneOf const labels")),
            "{:?}",
            proposal.warnings
        );
    }
}