# workflow_engine/core/values/schema.py
"""
Pydantic models that cover the subset of the JSON Schema language generated by
Pydantic for Value subclasses.
"""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from functools import cached_property
from typing import Any, ClassVar, Literal

from overrides import override
from pydantic import (
    ConfigDict,
    Field,
    ValidationError,
    model_serializer,
    SerializerFunctionWrapHandler,
    model_validator,
)

from ...utils.immutable import ImmutableBaseModel
from .data import Data, DataValue, build_data_type
from .mapping import StringMapValue
from .primitives import (
    BooleanValue,
    FloatValue,
    IntegerValue,
    NullValue,
    StringValue,
)
from .sequence import SequenceValue
from .value import Value, ValueType, value_type_registry


def merge_defs(
    a: Mapping[str, ValueSchema] | None,
    b: Mapping[str, ValueSchema] | None,
) -> Mapping[str, ValueSchema] | None:
    if a is None:
        return {} if b is None else b
    if b is None:
        return {} if a is None else a
    return dict(**a, **b)


class BaseValueSchema(ImmutableBaseModel):
    # We will only handle fields that Pydantic actually uses.
    # extra="forbid" offensively validates that we haven't missed any fields.
    model_config: ClassVar[ConfigDict] = ConfigDict(
        extra="forbid",
        serialize_by_alias=True,
    )

    # NOTE: recursive validation does not work unless this is a dict instead of a Mapping
    defs: dict[str, ValueSchema] = Field(
        default_factory=dict,
        serialization_alias="$defs",
        validation_alias="$defs",
    )
    title: str | None = None
    description: str | None = None
    default: Any | None = None
    const: Any | None = None

    @model_validator(mode="before")
    @classmethod
    def rename_defs_alias(cls, data: Any) -> Any:
        """
        A Pydantic bug causes model_validate() to break when called on an
        instance of BaseValueSchema.
        This is because the validator receives the .defs property from the
        object, whereas it is expecting $defs.
        To fix this, we move the property to the expected key.
        """
        if isinstance(data, Mapping) and ("$defs" not in data) and ("defs" in data):
            data = dict(data)
            data["$defs"] = data.pop("defs")
        return data

    @model_serializer(mode="wrap")
    def serialize_without_unset(self, handler: SerializerFunctionWrapHandler):
        """
        Custom serializer that excludes unset fields by default, so we don't
        clutter the output with unused fields.
        Applies to all subclasses of BaseValueSchema.

        Newbies might do this by overriding .model_dump(), however this is a
        trap because .model_dump() is not called recursively when this model is
        nested inside another model.

        Has special handling for $defs (included as long as it is not empty)
        and $ref (always included).
        """
        # Get the serialized data from the handler
        data: dict[str, Any] = handler(self)

        keys_to_keep: set[str] = {"$ref"}  # always keep $ref if present

        if hasattr(self, "__pydantic_fields_set__"):
            keys_to_keep.update(self.__pydantic_fields_set__)

        # special handling for $defs, because aliasing breaks
        # __pydantic_fields_set__
        if "$defs" in data and not (
            isinstance(data["$defs"], Mapping) and len(data["$defs"]) == 0
        ):
            keys_to_keep.add("$defs")

        return {k: v for k, v in data.items() if k in keys_to_keep}

    def to_value_cls(
        self,
        *extra_defs: Mapping[str, ValueSchema],
    ) -> ValueType:
        """
        Resolves this schema to a Pydantic class.
        If the title matches a concrete Value class, then it returns that class
        right away, regardless the other contents of the schema.
        This is technically a hack, but it is a massive quality of life boost
        for users, who can just provide { "title": "CustomValue" } as an alias
        for the entire CustomValue schema.

        References, if any, are resolved using self.defs first, then any
        extra_defs in order of decreasing precedence.
        """
        if self.title is not None and self.title in value_type_registry:
            return value_type_registry[self.title]
        return self.build_value_cls(*extra_defs)

    def build_value_cls(
        self,
        *extra_defs: Mapping[str, ValueSchema],
    ) -> ValueType:
        """
        Builds a Pydantic class from this schema. References, if any, are
        resolved using self.defs first, then any extra_defs in order of decreasing precedence.
        """
        raise NotImplementedError("Subclasses must implement this method")


class BooleanValueSchema(BaseValueSchema):
    type: Literal["boolean"]

    @override
    def build_value_cls(
        self,
        *extra_defs: Mapping[str, ValueSchema],
    ) -> type[BooleanValue]:
        return BooleanValue


class NullValueSchema(BaseValueSchema):
    type: Literal["null"]

    @override
    def build_value_cls(
        self,
        *extra_defs: Mapping[str, ValueSchema],
    ) -> type[NullValue]:
        return NullValue


class IntegerValueSchema(BaseValueSchema):
    # unused JSON Schema fields include minimum, maximum, multipleOf, etc.
    type: Literal["integer"]

    @override
    def build_value_cls(
        self,
        *extra_defs: Mapping[str, ValueSchema],
    ) -> type[IntegerValue]:
        return IntegerValue


class FloatValueSchema(BaseValueSchema):
    # unused JSON Schema fields include minimum, maximum, multipleOf, etc.
    type: Literal["number"]

    @override
    def build_value_cls(
        self,
        *extra_defs: Mapping[str, ValueSchema],
    ) -> type[FloatValue]:
        return FloatValue


class StringValueSchema(BaseValueSchema):
    # unused JSON Schema fields include length
    type: Literal["string"]
    enum: Sequence[str] | None = None
    pattern: str | None = None

    @override
    def build_value_cls(
        self,
        *extra_defs: Mapping[str, ValueSchema],
    ) -> type[StringValue]:
        return StringValue


class SequenceValueSchema(BaseValueSchema):
    type: Literal["array"]
    items: ValueSchema

    @override
    def build_value_cls(
        self,
        *extra_defs: Mapping[str, ValueSchema],
    ) -> type[SequenceValue]:
        T = self.items.to_value_cls(self.defs, *extra_defs)
        return SequenceValue[T]


class StringMapValueSchema(BaseValueSchema):
    type: Literal["object"]
    additionalProperties: ValueSchema | Literal[True] = True

    @override
    def build_value_cls(
        self,
        *extra_defs: Mapping[str, ValueSchema],
    ) -> type[StringMapValue]:
        if self.additionalProperties is True:
            return StringMapValue[Value]
        else:
            V = self.additionalProperties.to_value_cls(self.defs, *extra_defs)
            return StringMapValue[V]


class DataValueSchema(BaseValueSchema):
    """
    Matches a DataValue[T] schema, for some class T that inherits from Data.
    """

    type: Literal["object"]

    # NOTE: recursive validation does not work unless this is a dict instead of a Mapping
    properties: dict[str, ValueSchema]
    additionalProperties: ValueSchema | bool = False
    required: Sequence[str] = Field(default_factory=list)

    def build_data_cls(
        self,
        *extra_defs: Mapping[str, ValueSchema],
    ) -> type[Data]:
        """
        Rebuilds the Data type T from the schema.
        This rebuilt Data type will not be the same as the original Data type T
        (i.e. equality checks will fail), but it will have the same fields.
        It will not necessarily have the same to_value_schema() as T, because
        the rebuilding process does not yet capture all of the descriptions from
        the original schema.
        """
        # TODO: capture the descriptions from the schema
        required_set = frozenset(self.required)
        properties = {
            k: (
                v.to_value_cls(self.defs, *extra_defs),
                k in required_set,
            )
            for k, v in self.properties.items()
        }
        assert self.title is not None
        return build_data_type(self.title, properties)

    @override
    def build_value_cls(
        self,
        *extra_defs: Mapping[str, ValueSchema],
    ) -> type[DataValue]:
        D = self.build_data_cls(*extra_defs)
        return DataValue[D]


class UnionValueSchema(BaseValueSchema):
    anyOf: Sequence[ValueSchema]

    @override
    def build_value_cls(
        self,
        *extra_defs: Mapping[str, ValueSchema],
    ) -> ValueType:
        # TODO: add proper type union support
        raise NotImplementedError(
            "There is no Value type handling for type unions yet."
        )


class ReferenceValueSchema(BaseValueSchema):
    """
    A schema of the form
    { "$defs": { "foo": ... }, "$ref": "#/$defs/foo" }
    """

    model_config: ClassVar[ConfigDict] = ConfigDict(
        serialize_by_alias=True,  # already set in BaseValueSchema, but setting again here to be explicit
    )

    ref: str = Field(
        pattern=r"^#/\$defs/\w+$",
        serialization_alias="$ref",
        validation_alias="$ref",
    )

    @model_validator(mode="before")
    @classmethod
    def rename_ref_alias(cls, data: Any) -> Any:
        """
        A Pydantic bug causes model_validate() to break when called on an
        instance of ReferenceValueSchema.
        This is because the validator receives the .ref property from the
        object, whereas it is expecting $ref.
        To fix this, we move the property to the expected key.
        """
        if isinstance(data, Mapping) and ("$ref" not in data) and ("ref" in data):
            data = dict(data)
            data["$ref"] = data.pop("ref")
            return data
        return data

    @cached_property
    def id(self) -> str:
        _hash, _defs, id = self.ref.split("/")
        assert _hash == "#"
        assert _defs == "$defs"
        return id

    @override
    def build_value_cls(
        self,
        *extra_defs: Mapping[str, ValueSchema],
    ) -> ValueType:
        for defs in (self.defs, *extra_defs):
            if self.id in defs:
                return defs[self.id].to_value_cls(self.defs, *extra_defs)
        raise KeyError(f"Schema definition for {self.id} not found")


type ValueSchema = (
    BooleanValueSchema
    | DataValueSchema
    | FloatValueSchema
    | IntegerValueSchema
    | NullValueSchema
    | SequenceValueSchema
    | StringMapValueSchema
    | StringValueSchema
    | UnionValueSchema
    | ReferenceValueSchema  # must be handled second-to-last
    | BaseValueSchema  # must be handled last, corresponds to the Any catch-all
)


# yep, this is possible
class ValueSchemaValue(Value[ValueSchema]):
    pass


def validate_value_schema(schema: Any) -> ValueSchema:
    # idempotent: if the schema is already a ValueSchema, just return it
    if isinstance(schema, BaseValueSchema):
        return schema
    try:
        return ValueSchemaValue.model_validate(schema).root
    except ValidationError as e:
        raise ValueError(f"Invalid value schema: {schema}") from e


__all__ = [
    "ValueSchema",
    "ValueSchemaValue",
]
