genson-core 1.0.3

use serde_json::{json, Value};
use std::hash::Hasher;

use crate::dag::{kw, Dag, Id, ObjBuf};
use crate::genson_rs::interned::{RowStates, StateId};
use crate::genson_rs::node::{DataType, Deferred, SchemaNode};

const DEFAULT_SCHEMA_URI: &str = "http://json-schema.org/schema#";
const NULL_SCHEMA_URI: &str = "NULL";

pub struct SchemaBuilder {
    schema_uri: Option<String>,
    root_node: SchemaNode,
}

impl SchemaBuilder {
    /// Create a new SchemaBuilder object. The schema_uri parameter is optional, a value
    /// of "AUTO" will automatically detect the schema URI based on the input schema, if no
    /// schema URI was detected, a default URI of "<http://json-schema.org/schema#>" will be used.
    /// A value of None will leave out the "$schema" keyword in the output schema.
    pub fn new(schema_uri: Option<&str>) -> Self {
        // TODO: the functionality to allow non-default node class with extended
        //  strategies is not supported yet
        let root_node: SchemaNode = SchemaNode::new();

        if let Some(uri) = schema_uri {
            if uri == "AUTO" {
                SchemaBuilder {
                    schema_uri: None,
                    root_node,
                }
            } else {
                SchemaBuilder {
                    schema_uri: Some(uri.to_string()),
                    root_node,
                }
            }
        } else {
            SchemaBuilder {
                schema_uri: Some(NULL_SCHEMA_URI.to_string()),
                root_node,
            }
        }
    }

    /// Merge in raw JSON schema object
    pub fn add_schema(&mut self, mut schema: Value) {
        self.add_schema_mut(&mut schema);
    }

    /// Like [`add_schema`](Self::add_schema) but leaves `schema` with the caller, so it can be
    /// dropped elsewhere (a large schema is costly to free). Removes `$schema` from it when
    /// that is adopted as this builder's schema URI.
    pub fn add_schema_mut(&mut self, schema: &mut Value) {
        if let Value::Object(schema_obj) = schema {
            if schema_obj.contains_key("$schema") && self.schema_uri.is_none() {
                self.schema_uri = Some(schema_obj["$schema"].as_str().unwrap().to_string());
                schema_obj.shift_remove("$schema");
            }
            self.root_node.add_schema(DataType::Schema(schema));
        } else {
            panic!("Invalid schema type - must be a valid JSON object")
        }
    }

    /// Merge `schemas` in order, same result as calling [`add_schema_mut`](Self::add_schema_mut)
    /// on each, with the merge parallelised across the root object's properties.
    pub fn add_schemas_mut(&mut self, schemas: &mut [Value]) {
        for schema in schemas.iter_mut() {
            let Value::Object(schema_obj) = schema else {
                panic!("Invalid schema type - must be a valid JSON object")
            };
            if schema_obj.contains_key("$schema") && self.schema_uri.is_none() {
                self.schema_uri = Some(schema_obj["$schema"].as_str().unwrap().to_string());
                schema_obj.shift_remove("$schema");
            }
        }
        let refs: Vec<&Value> = schemas.iter().collect();
        if !self.root_node.add_object_schemas_par(&refs) {
            for schema in refs {
                self.root_node.add_schema(DataType::Schema(schema));
            }
        }
    }

    /// Merge another builder into this one. When this builder is still empty and both
    /// share the same `$schema` handling, takes over its nodes instead of round-tripping
    /// through a schema value.
    pub fn absorb(&mut self, other: SchemaBuilder) {
        // `other` with the NULL uri emits no "$schema", so nothing else would be merged.
        if self.root_node.is_empty() && other.schema_uri.as_deref() == Some(NULL_SCHEMA_URI) {
            self.root_node = other.root_node;
        } else {
            self.add_builder(other);
        }
    }

    /// Merge another builder into this one, with the same result as
    /// `self.add_schema(other.to_schema())` but merging its nodes directly.
    pub fn add_builder(&mut self, other: SchemaBuilder) {
        match other.emitted_schema_uri() {
            None => self.root_node.absorb(other.root_node),
            // `add_schema_mut` adopts the `$schema` and merges the rest
            Some(uri) if self.schema_uri.is_none() => {
                self.schema_uri = Some(uri.to_string());
                self.root_node.absorb(other.root_node);
            }
            // Otherwise the `$schema` keyword itself is merged into the root
            Some(_) => self.add_schema(other.to_schema()),
        }
    }

    /// Merge `others` in order, with the same result as `add_builder` on each in turn but
    /// with the roots that are objects merged in parallel across their properties.
    pub fn add_builders(&mut self, others: Vec<SchemaBuilder>) {
        let mut objects = Vec::with_capacity(others.len());
        let mut kept_uri = None;
        for other in others {
            if !other.root_node.is_object() {
                self.root_node.absorb_all(std::mem::take(&mut objects));
                self.add_builder(other);
                continue;
            }
            match (other.emitted_schema_uri(), &self.schema_uri) {
                (Some(uri), None) => self.schema_uri = Some(uri.to_string()),
                // `add_schema` merges an unadopted `$schema` into the root as a keyword
                (Some(uri), Some(_)) => kept_uri = kept_uri.or(Some(uri.to_string())),
                (None, _) => {}
            }
            objects.push(other.root_node);
        }
        self.root_node.absorb_all(objects);
        if let Some(uri) = kept_uri {
            self.root_node.add_object_keyword("$schema", uri.into());
        }
    }

    /// `add_builders` of each chunk of the rows' builders that interned row states stand
    /// for, in turn (each a builder made with `row_uri`'s setting, its root an object): the
    /// same result, read from the states. The root takes every row (and each chunk's
    /// `$schema`) in turn, and its children their rows' states after the last chunk.
    pub(crate) fn add_interned_row_chunks<'a>(
        &mut self,
        chunks: impl IntoIterator<Item = Vec<(&'a RowStates, StateId)>>,
        row_uri: Option<&str>,
    ) {
        let row_uri = SchemaBuilder::new(row_uri)
            .emitted_schema_uri()
            .map(str::to_string);
        let mut deferred = Deferred::default();
        for chunk in chunks {
            let mut kept_uri = None;
            for (states, state) in chunk {
                match (&row_uri, &self.schema_uri) {
                    (Some(uri), None) => self.schema_uri = Some(uri.clone()),
                    // `add_schema` merges an unadopted `$schema` into the root as a keyword
                    (Some(uri), Some(_)) => kept_uri = kept_uri.or(Some(uri.clone())),
                    (None, _) => {}
                }
                self.root_node
                    .absorb_interned_deferred(states, state, &mut deferred);
            }
            if let Some(uri) = kept_uri {
                self.root_node.add_object_keyword("$schema", uri.into());
            }
        }
        self.root_node.absorb_deferred(deferred);
    }

    pub(crate) fn root_node_mut(&mut self) -> &mut SchemaNode {
        &mut self.root_node
    }

    /// A hash of the schema `to_schema` gives, computed without building it: builders with
    /// equal schemas hash equal.
    pub fn schema_hash(&self) -> u64 {
        let mut hasher = xxhash_rust::xxh64::Xxh64::new(0);
        hasher.write(self.emitted_schema_uri().unwrap_or_default().as_bytes());
        hasher.write_u8(0xff);
        self.root_node.hash_schema(&mut hasher);
        hasher.finish()
    }

    /// The `$schema` that `to_schema` writes, if any.
    fn emitted_schema_uri(&self) -> Option<&str> {
        match self.schema_uri.as_deref() {
            Some(NULL_SCHEMA_URI) => None,
            Some(uri) => Some(uri),
            None => Some(DEFAULT_SCHEMA_URI),
        }
    }

    /// Merge multiple raw JSON schema objects at once
    pub fn add_schemas(&mut self, schemas: &[Value]) {
        let schemas: Vec<&Value> = schemas.iter().collect();
        self.root_node.add_schemas(&schemas);
    }

    /// Merge in another SchemaNode object
    pub fn add_schema_node(&mut self, node: SchemaNode) {
        self.root_node.add_schema(DataType::SchemaNode(&node));
    }

    /// Modify the schema to accomodate the input object
    pub fn add_object(&mut self, object: super::JsonValue) {
        self.root_node.add_object(DataType::Object(object));
    }

    /// Export the currently constructed schema as a JSON object
    pub fn to_schema(&self) -> Value {
        let mut base_schema = self.get_base_schema();

        let base_schema_map = base_schema.as_object_mut().unwrap();
        let Value::Object(node_schema_map) = self.root_node.to_schema() else {
            panic!("node schema must be an object")
        };

        for (key, value) in node_schema_map {
            base_schema_map.insert(key, value);
        }
        base_schema
    }

    /// `to_schema` as a DAG.
    pub(crate) fn to_dag(&self, dag: &mut Dag) -> Id {
        let mut schema = ObjBuf::default();
        match self.schema_uri.as_deref() {
            Some(NULL_SCHEMA_URI) => {}
            uri => {
                let uri = dag.str(uri.unwrap_or(DEFAULT_SCHEMA_URI));
                schema.insert(kw::SCHEMA, uri);
            }
        }
        let node = self.root_node.to_dag(dag);
        let Some(node) = dag.as_object(node) else {
            panic!("node schema must be an object")
        };
        for (key, value) in node.iter() {
            schema.insert(key, value);
        }
        schema.build(dag)
    }

    fn get_base_schema(&self) -> Value {
        if let Some(uri) = &self.schema_uri {
            if uri == NULL_SCHEMA_URI {
                json!({})
            } else {
                json!({"$schema": uri})
            }
        } else {
            json!({"$schema": DEFAULT_SCHEMA_URI})
        }
    }

    /// Serialize the currently constructed schema to a JSON string
    pub fn to_json(&self) -> String {
        let schema = self.to_schema();
        schema.to_string()
    }
}