ploidy-core 0.15.0

use std::num::NonZeroUsize;

use itertools::Itertools;
use rustc_hash::FxHashMap;

use crate::{
    arena::Arena,
    ir::{JsonF64, SchemaTypeInfo},
    parse::{AdditionalProperties, Document, Format, RefOrSchema, Schema, Ty},
};

use super::types::{
    Enum, EnumVariant, InlineTypeId, InlineTypeIds, PrimitiveType, SpecContainer, SpecInlineType,
    SpecInner, SpecSchemaType, SpecStruct, SpecStructField, SpecTagged, SpecTaggedVariant,
    SpecType, SpecUntagged, StructFieldName,
};

/// Metadata about a type in the dependency graph.
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub enum TypeInfo<'a> {
    Schema(SchemaTypeInfo<'a>),
    Inline(InlineTypeId),
}

impl<'a> From<SchemaTypeInfo<'a>> for TypeInfo<'a> {
    fn from(info: SchemaTypeInfo<'a>) -> Self {
        Self::Schema(info)
    }
}

impl From<InlineTypeId> for TypeInfo<'_> {
    fn from(id: InlineTypeId) -> Self {
        Self::Inline(id)
    }
}

/// Context for the [`IrTransformer`].
#[derive(Debug)]
pub struct TransformContext<'a> {
    arena: &'a Arena,
    doc: &'a Document,
    ids: InlineTypeIds<'a>,
}

impl<'a> TransformContext<'a> {
    /// Creates a new context for the given document.
    pub fn new(arena: &'a Arena, doc: &'a Document, ids: InlineTypeIds<'a>) -> Self {
        Self { arena, doc, ids }
    }
}

pub(super) fn transform_with_context<'context, 'a>(
    context: &'context TransformContext<'a>,
    name: impl Into<TypeInfo<'a>>,
    schema: &'a Schema,
) -> SpecType<'a> {
    IrTransformer::new(context, name.into(), schema).transform()
}

#[derive(Debug)]
struct IrTransformer<'context, 'a> {
    context: &'context TransformContext<'a>,
    name: TypeInfo<'a>,
    schema: &'a Schema,
}

impl<'context, 'a> IrTransformer<'context, 'a> {
    fn new(
        context: &'context TransformContext<'a>,
        name: TypeInfo<'a>,
        schema: &'a Schema,
    ) -> Self {
        Self {
            context,
            name,
            schema,
        }
    }

    /// Returns a reference to the arena allocator.
    #[inline]
    fn arena(&self) -> &'a Arena {
        self.context.arena
    }

    fn transform(self) -> SpecType<'a> {
        self.try_tagged()
            .or_else(Self::try_untagged)
            .or_else(Self::try_any_of)
            .or_else(Self::try_enum)
            .or_else(Self::try_struct)
            .unwrap_or_else(Self::other)
    }

    fn try_tagged(self) -> Result<SpecType<'a>, Self> {
        let (Some(one_of), Some(discriminator)) = (&self.schema.one_of, &self.schema.discriminator)
        else {
            return Err(self);
        };

        let variants = {
            let mut inverted = FxHashMap::<_, Vec<_>>::default();
            for (tag, r) in &discriminator.mapping {
                inverted.entry(r).or_default().push(tag.as_str());
            }
            let mut variants = Vec::with_capacity(one_of.len());
            for schema in one_of {
                match schema {
                    RefOrSchema::Ref(r) => {
                        let name: &_ = self.arena().alloc_str(&r.name());
                        let aliases = match inverted.get(r).map(|s| s.as_slice()).unwrap_or(&[]) {
                            // When a discriminator value doesn't have
                            // an explicit `mapping`, use the schema name.
                            [] => &[name],
                            aliases => aliases,
                        };
                        variants.push(SpecTaggedVariant {
                            name,
                            ty: self.arena().alloc(SpecType::Ref(r)),
                            aliases: self.arena().alloc_slice_copy(aliases),
                        });
                    }
                    // An inline schema variant can't have a discriminator mapping;
                    // fall through to `try_untagged`.
                    RefOrSchema::Inline(_) => return Err(self),
                }
            }
            variants
        };

        let tagged = SpecTagged {
            description: self.schema.description.as_deref(),
            tag: discriminator.property_name.as_str(),
            variants: self.arena().alloc_slice_copy(&variants),
            fields: self.arena().alloc_slice(self.properties()),
            parents: self.arena().alloc_slice(self.parents()),
        };

        Ok(match self.name {
            TypeInfo::Schema(info) => SpecSchemaType::Tagged(info, tagged).into(),
            TypeInfo::Inline(id) => SpecInlineType::Tagged(id, tagged).into(),
        })
    }

    fn try_untagged(self) -> Result<SpecType<'a>, Self> {
        let Some(one_of) = &self.schema.one_of else {
            return Err(self);
        };

        let variants = match &**one_of {
            [] => {
                return Ok(match self.name {
                    TypeInfo::Schema(info) => SpecSchemaType::Any(info).into(),
                    TypeInfo::Inline(id) => SpecInlineType::Any(id).into(),
                });
            }
            [schema] => {
                // Unwrap single-variant untagged unions.
                return Ok(match schema {
                    RefOrSchema::Ref(r) => SpecType::Ref(r),
                    RefOrSchema::Inline(s) if matches!(&*s.ty, [Ty::Null]) => match self.name {
                        TypeInfo::Schema(info) => SpecSchemaType::Any(info).into(),
                        TypeInfo::Inline(id) => SpecInlineType::Any(id).into(),
                    },
                    RefOrSchema::Inline(schema) => {
                        transform_with_context(self.context, self.name, schema)
                    }
                });
            }
            variants => variants
                .iter()
                .map(|schema| {
                    let ty = match schema {
                        RefOrSchema::Ref(r) => Some(SpecType::Ref(r)),
                        RefOrSchema::Inline(s) if matches!(&*s.ty, [Ty::Null]) => None,
                        RefOrSchema::Inline(schema) => {
                            let id = self.context.ids.next();
                            Some(transform_with_context(self.context, id, schema))
                        }
                    };
                    ty.map(|ty| &*self.arena().alloc(ty))
                })
                .collect_vec(),
        };

        Ok(match &*variants {
            // Simplify two-variant untagged unions, where one is a type
            // and the other is `null`, into optionals.
            [Some(ty), None] | [None, Some(ty)] => {
                let container = SpecContainer::Optional(SpecInner {
                    description: self.schema.description.as_deref(),
                    ty,
                });
                match self.name {
                    TypeInfo::Schema(info) => SpecSchemaType::Container(info, container).into(),
                    TypeInfo::Inline(id) => SpecInlineType::Container(id, container).into(),
                }
            }

            variants => {
                let untagged = SpecUntagged {
                    description: self.schema.description.as_deref(),
                    variants: self.arena().alloc_slice_copy(variants),
                    fields: self.arena().alloc_slice(self.properties()),
                    parents: self.arena().alloc_slice(self.parents()),
                };
                match self.name {
                    TypeInfo::Schema(info) => SpecSchemaType::Untagged(info, untagged).into(),
                    TypeInfo::Inline(id) => SpecInlineType::Untagged(id, untagged).into(),
                }
            }
        })
    }

    fn try_any_of(self) -> Result<SpecType<'a>, Self> {
        let Some(any_of) = &self.schema.any_of else {
            return Err(self);
        };
        if let [schema] = &**any_of {
            // A single-variant `anyOf` should unwrap to the variant type. This
            // preserves type references that would otherwise become `Any`.
            return Ok(match schema {
                RefOrSchema::Ref(r) => SpecType::Ref(r),
                RefOrSchema::Inline(schema) => {
                    transform_with_context(self.context, self.name, schema)
                }
            });
        }

        let any_of_fields = any_of
            .iter()
            .enumerate()
            .map(|(index, schema)| {
                let ordinal = NonZeroUsize::new(index + 1).unwrap();
                let (field_name, ty, description) = match schema {
                    RefOrSchema::Ref(r) => {
                        let name = StructFieldName::Name(self.arena().alloc_str(&r.name()));
                        let ty: &_ = self.arena().alloc(SpecType::Ref(r));
                        let desc = r
                            .pointer()
                            .follow::<&Schema>(self.context.doc)
                            .ok()
                            .and_then(|s| s.description.as_deref());
                        (name, ty, desc)
                    }
                    RefOrSchema::Inline(schema) => {
                        let name = StructFieldName::Ordinal(ordinal);
                        let id = self.context.ids.next();
                        let ty: &_ =
                            self.arena()
                                .alloc(transform_with_context(self.context, id, schema));
                        let desc = schema.description.as_deref();
                        (name, ty, desc)
                    }
                };
                // Flattened `anyOf` fields are always optional.
                let id = self.context.ids.next();
                let ty: &_ = self.arena().alloc(
                    SpecInlineType::Container(
                        id,
                        SpecContainer::Optional(SpecInner { description, ty }),
                    )
                    .into(),
                );
                SpecStructField {
                    name: field_name,
                    ty,
                    required: false,
                    description,
                    flattened: true,
                }
            })
            .collect_vec();

        let ty = SpecStruct {
            description: self.schema.description.as_deref(),
            fields: self.arena().alloc_slice({
                // Combine all the fields: regular properties first,
                // followed by the flattened `anyOf` fields. This ordering
                // ensures that regular properties take precedence during
                // (de)serialization.
                itertools::chain!(self.properties(), any_of_fields)
            }),
            parents: self.arena().alloc_slice(self.parents()),
        };

        Ok(match self.name {
            TypeInfo::Schema(info) => SpecSchemaType::Struct(info, ty).into(),
            TypeInfo::Inline(id) => SpecInlineType::Struct(id, ty).into(),
        })
    }

    fn try_enum(self) -> Result<SpecType<'a>, Self> {
        let Some(values) = &self.schema.variants else {
            return Err(self);
        };
        // JSON Schema Validation (draft-bhutton-json-schema-validation-01)
        // recommends unique enum values, but specs in the wild repeat values.
        let variants = self.arena().alloc_slice(
            values
                .iter()
                .filter_map(|value| {
                    if let Some(s) = value.as_str() {
                        Some(EnumVariant::String(s))
                    } else if let Some(n) = value.as_number() {
                        if let Some(n) = n.as_i64() {
                            Some(EnumVariant::I64(n))
                        } else if let Some(n) = n.as_u64() {
                            Some(EnumVariant::U64(n))
                        } else {
                            n.as_f64().map(|f| EnumVariant::F64(JsonF64::new(f)))
                        }
                    } else {
                        value.as_bool().map(EnumVariant::Bool)
                    }
                })
                .unique(),
        );
        let ty = Enum {
            description: self.schema.description.as_deref(),
            variants,
        };
        Ok(match self.name {
            TypeInfo::Schema(info) => SpecSchemaType::Enum(info, ty).into(),
            TypeInfo::Inline(id) => SpecInlineType::Enum(id, ty).into(),
        })
    }

    fn try_struct(self) -> Result<SpecType<'a>, Self> {
        if self.schema.properties.is_none() && self.schema.all_of.is_none() {
            return Err(self);
        }

        let ty = SpecStruct {
            description: self.schema.description.as_deref(),
            fields: self.arena().alloc_slice(itertools::chain!(
                self.properties(),
                self.additional_properties()
            )),
            parents: self.arena().alloc_slice(self.parents()),
        };
        Ok(match self.name {
            TypeInfo::Schema(info) => SpecSchemaType::Struct(info, ty).into(),
            TypeInfo::Inline(id) => SpecInlineType::Struct(id, ty).into(),
        })
    }

    fn other(self) -> SpecType<'a> {
        let mut other = Other {
            variants: Vec::with_capacity(self.schema.ty.len()),
            nullable: false,
        };

        for ty in &self.schema.ty {
            let variant = match (ty, self.schema.format) {
                (Ty::String, Some(Format::DateTime)) => {
                    OtherVariant::Primitive(PrimitiveType::DateTime)
                }
                (Ty::String, Some(Format::Date)) => OtherVariant::Primitive(PrimitiveType::Date),
                (Ty::String, Some(Format::Uri)) => OtherVariant::Primitive(PrimitiveType::Url),
                (Ty::String, Some(Format::Uuid)) => OtherVariant::Primitive(PrimitiveType::Uuid),
                (Ty::String, Some(Format::Byte)) => OtherVariant::Primitive(PrimitiveType::Bytes),
                (Ty::String, Some(Format::Binary)) => {
                    OtherVariant::Primitive(PrimitiveType::Binary)
                }
                (Ty::String, _) => OtherVariant::Primitive(PrimitiveType::String),

                (Ty::Integer, Some(Format::Int8)) => OtherVariant::Primitive(PrimitiveType::I8),
                (Ty::Integer, Some(Format::UInt8)) => OtherVariant::Primitive(PrimitiveType::U8),
                (Ty::Integer, Some(Format::Int16)) => OtherVariant::Primitive(PrimitiveType::I16),
                (Ty::Integer, Some(Format::UInt16)) => OtherVariant::Primitive(PrimitiveType::U16),
                (Ty::Integer, Some(Format::Int32)) => OtherVariant::Primitive(PrimitiveType::I32),
                (Ty::Integer, Some(Format::UInt32)) => OtherVariant::Primitive(PrimitiveType::U32),
                (Ty::Integer, Some(Format::Int64)) => OtherVariant::Primitive(PrimitiveType::I64),
                (Ty::Integer, Some(Format::UInt64)) => OtherVariant::Primitive(PrimitiveType::U64),
                (Ty::Integer, Some(Format::UnixTime)) => {
                    OtherVariant::Primitive(PrimitiveType::UnixTime)
                }
                (Ty::Integer, _) => OtherVariant::Primitive(PrimitiveType::I32),

                (Ty::Number, Some(Format::Float)) => OtherVariant::Primitive(PrimitiveType::F32),
                (Ty::Number, Some(Format::Double)) => OtherVariant::Primitive(PrimitiveType::F64),
                (Ty::Number, Some(Format::UnixTime)) => {
                    OtherVariant::Primitive(PrimitiveType::UnixTime)
                }
                (Ty::Number, _) => OtherVariant::Primitive(PrimitiveType::F64),

                (Ty::Boolean, _) => OtherVariant::Primitive(PrimitiveType::Bool),

                (Ty::Array, _) => {
                    let items = match &self.schema.items {
                        Some(RefOrSchema::Ref(r)) => SpecType::Ref(r),
                        Some(RefOrSchema::Inline(schema)) => {
                            let id = self.context.ids.next();
                            transform_with_context(self.context, id, schema)
                        }
                        None => {
                            let id = self.context.ids.next();
                            SpecInlineType::Any(id).into()
                        }
                    };
                    OtherVariant::Array(SpecInner {
                        description: self.schema.description.as_deref(),
                        ty: self.arena().alloc(items),
                    })
                }

                (Ty::Object, _) => {
                    let inner = match &self.schema.additional_properties {
                        Some(AdditionalProperties::RefOrSchema(RefOrSchema::Ref(r))) => {
                            Some(SpecInner {
                                description: self.schema.description.as_deref(),
                                ty: self.arena().alloc(SpecType::Ref(r)),
                            })
                        }
                        Some(AdditionalProperties::RefOrSchema(RefOrSchema::Inline(schema))) => {
                            let id = self.context.ids.next();
                            Some(SpecInner {
                                description: self.schema.description.as_deref(),
                                ty: self.arena().alloc(transform_with_context(
                                    self.context,
                                    id,
                                    schema,
                                )),
                            })
                        }
                        Some(AdditionalProperties::Bool(true)) => {
                            let id = self.context.ids.next();
                            Some(SpecInner {
                                description: self.schema.description.as_deref(),
                                ty: self
                                    .arena()
                                    .alloc(SpecType::Inline(SpecInlineType::Any(id))),
                            })
                        }
                        _ => None,
                    };
                    match inner {
                        Some(inner) => OtherVariant::Map(inner),
                        None => OtherVariant::Any,
                    }
                }

                (Ty::Null, _) => {
                    other.nullable = true;
                    continue;
                }
            };
            other.variants.push(variant);
        }

        match (&*other.variants, other.nullable) {
            // An empty `type` array is invalid in JSON Schema,
            // but we treat it as "any type".
            ([], false) => match self.name {
                TypeInfo::Schema(info) => SpecSchemaType::Any(info).into(),
                TypeInfo::Inline(id) => SpecInlineType::Any(id).into(),
            },

            // A `null` variant becomes `Any`.
            ([], true) => match self.name {
                TypeInfo::Schema(info) => SpecSchemaType::Any(info).into(),
                TypeInfo::Inline(id) => SpecInlineType::Any(id).into(),
            },

            // A union with a single, non-`null` variant unwraps to
            // the type of that variant.
            ([variant], false) => match self.name {
                TypeInfo::Schema(info) => variant.to_schema_type(info).into(),
                TypeInfo::Inline(id) => variant.to_inline_type(id).into(),
            },

            // A two-variant union, with one type T and one `null` variant,
            // simplifies to `Optional(T)`.
            ([variant], true) => {
                let container = SpecContainer::Optional(SpecInner {
                    description: self.schema.description.as_deref(),
                    ty: self
                        .arena()
                        .alloc(variant.to_inline_type(self.context.ids.next()).into()),
                });
                match self.name {
                    TypeInfo::Schema(info) => SpecSchemaType::Container(info, container).into(),
                    TypeInfo::Inline(id) => SpecInlineType::Container(id, container).into(),
                }
            }

            // Anything else becomes an untagged union.
            (many, nullable) => {
                let mut variants = many
                    .iter()
                    .map(|variant| {
                        Some(
                            &*self
                                .arena()
                                .alloc(variant.to_inline_type(self.context.ids.next()).into()),
                        )
                    })
                    .collect_vec();
                if nullable {
                    variants.push(None);
                }
                let untagged = SpecUntagged {
                    description: self.schema.description.as_deref(),
                    variants: self.arena().alloc_slice_copy(&variants),
                    fields: &[],
                    parents: self.arena().alloc_slice(self.parents()),
                };
                match self.name {
                    TypeInfo::Schema(info) => SpecSchemaType::Untagged(info, untagged).into(),
                    TypeInfo::Inline(id) => SpecInlineType::Untagged(id, untagged).into(),
                }
            }
        }
    }

    // MARK: Shared lowering

    /// Lowers immediate parents from `allOf` into a list of types.
    fn parents(&self) -> impl Iterator<Item = &'a SpecType<'a>> {
        self.schema
            .all_of
            .iter()
            .flatten()
            .enumerate()
            .map(|(index, parent)| (index + 1, parent))
            .map(move |(_index, parent)| &*match parent {
                RefOrSchema::Ref(r) => self.arena().alloc(SpecType::Ref(r)),
                RefOrSchema::Inline(schema) => {
                    let id = self.context.ids.next();
                    self.arena()
                        .alloc(transform_with_context(self.context, id, schema))
                }
            })
    }

    /// Lowers regular fields from `properties` into struct fields,
    /// wrapping nullable and optional fields in [`Container::Optional`].
    fn properties(&self) -> impl Iterator<Item = SpecStructField<'a>> {
        self.schema
            .properties
            .iter()
            .flatten()
            .map(move |(name, field_schema)| {
                let field_name = name.as_str();
                let required = self.schema.required.contains(name);
                let ty: &_ = match field_schema {
                    RefOrSchema::Ref(r) => self.arena().alloc(SpecType::Ref(r)),
                    RefOrSchema::Inline(schema) => {
                        let id = self.context.ids.next();
                        self.arena()
                            .alloc(transform_with_context(self.context, id, schema))
                    }
                };
                let description = match field_schema {
                    RefOrSchema::Inline(schema) => schema.description.as_deref(),
                    RefOrSchema::Ref(r) => r
                        .pointer()
                        .follow::<&Schema>(self.context.doc)
                        .ok()
                        .and_then(|schema| schema.description.as_deref()),
                };
                let nullable = match field_schema {
                    RefOrSchema::Inline(schema) if schema.nullable => true,
                    RefOrSchema::Ref(r) => r
                        .pointer()
                        .follow::<&Schema>(self.context.doc)
                        .is_ok_and(|schema| schema.nullable),
                    _ => false,
                };
                // Wrap the type in `Optional` if the field is either
                // explicitly nullable, or implicitly optional. The `required`
                // flag distinguishes between the two for codegen.
                let ty: &_ = if nullable || !required {
                    let id = self.context.ids.next();
                    self.arena().alloc(SpecType::from(SpecInlineType::Container(
                        id,
                        SpecContainer::Optional(SpecInner { description, ty }),
                    )))
                } else {
                    ty
                };
                SpecStructField {
                    name: StructFieldName::Name(field_name),
                    ty,
                    required,
                    description,
                    flattened: false,
                }
            })
    }

    /// Lowers `additionalProperties` into a struct field definition,
    /// if the schema specifies them.
    fn additional_properties(&self) -> Option<SpecStructField<'a>> {
        let inner = match &self.schema.additional_properties {
            Some(AdditionalProperties::RefOrSchema(RefOrSchema::Ref(r))) => SpecInner {
                description: self.schema.description.as_deref(),
                ty: self.arena().alloc(SpecType::Ref(r)),
            },
            Some(AdditionalProperties::RefOrSchema(RefOrSchema::Inline(schema))) => {
                let id = self.context.ids.next();
                SpecInner {
                    description: self.schema.description.as_deref(),
                    ty: self
                        .arena()
                        .alloc(transform_with_context(self.context, id, schema)),
                }
            }
            Some(AdditionalProperties::Bool(true)) => {
                let id = self.context.ids.next();
                SpecInner {
                    description: self.schema.description.as_deref(),
                    ty: self
                        .arena()
                        .alloc(SpecType::Inline(SpecInlineType::Any(id))),
                }
            }
            _ => return None,
        };

        let map_id = self.context.ids.next();
        let ty: &_ = self.arena().alloc(SpecType::from(SpecInlineType::Container(
            map_id,
            SpecContainer::Map(inner),
        )));

        Some(SpecStructField {
            name: StructFieldName::AdditionalProperties,
            ty,
            required: true,
            description: None,
            flattened: true,
        })
    }
}

/// A union of variants for representing OpenAPI 3.1-style
/// `type` arrays.
struct Other<'a> {
    variants: Vec<OtherVariant<'a>>,
    nullable: bool,
}

/// A variant of an [`Other`] union.
#[derive(Clone, Copy)]
enum OtherVariant<'a> {
    Primitive(PrimitiveType),
    Array(SpecInner<'a>),
    Map(SpecInner<'a>),
    Any,
}

impl<'a> OtherVariant<'a> {
    /// Converts this variant to a [`SpecSchemaType`].
    fn to_schema_type(self, info: SchemaTypeInfo<'a>) -> SpecSchemaType<'a> {
        match self {
            Self::Primitive(p) => SpecSchemaType::Primitive(info, p),
            Self::Array(inner) => SpecSchemaType::Container(info, SpecContainer::Array(inner)),
            Self::Map(inner) => SpecSchemaType::Container(info, SpecContainer::Map(inner)),
            Self::Any => SpecSchemaType::Any(info),
        }
    }

    /// Converts this variant to a [`SpecInlineType`].
    fn to_inline_type(self, id: InlineTypeId) -> SpecInlineType<'a> {
        match self {
            Self::Primitive(p) => SpecInlineType::Primitive(id, p),
            Self::Array(inner) => SpecInlineType::Container(id, SpecContainer::Array(inner)),
            Self::Map(inner) => SpecInlineType::Container(id, SpecContainer::Map(inner)),
            Self::Any => SpecInlineType::Any(id),
        }
    }
}