# 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 Annotated, Any, ClassVar, Final, Literal, Self, TypeVar

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

from workflow_engine.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, ValueRegistry, ValueType
from workflow_engine.utils.hash import json_digest

_T_item = TypeVar("_T_item", bound=Value)


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)


# Maps JSON Schema keywords to pydantic Field kwargs for each value category.
_NUMERIC_FIELD_MAP: dict[str, str] = {
    "minimum": "ge",
    "maximum": "le",
    "exclusiveMinimum": "gt",
    "exclusiveMaximum": "lt",
    "multipleOf": "multiple_of",
}
_STRING_FIELD_MAP: dict[str, str] = {
    "minLength": "min_length",
    "maxLength": "max_length",
    "pattern": "pattern",
}
_SEQUENCE_FIELD_MAP: dict[str, str] = {
    "minItems": "min_length",
    "maxItems": "max_length",
}
_MAP_FIELD_MAP: dict[str, str] = {
    "minProperties": "min_length",
    "maxProperties": "max_length",
}


def _build_constrained_cls(
    base_cls: type,
    field_map: dict[str, str],
    extras: dict[str, Any],
) -> type:
    """
    Return a constrained subclass of *base_cls*, or *base_cls* itself if there
    is nothing to add.

    Known JSON Schema keywords (per *field_map*) become Pydantic Field
    constraints on the ``root`` annotation, so they are both enforced at
    runtime and reflected correctly in the schema.  Anything not in the map
    is passed through as ``json_schema_extra`` (schema-only, no enforcement).
    """
    if not extras:
        return base_cls

    field_kwargs: dict[str, Any] = {}
    schema_extras: dict[str, Any] = {}
    for key, value in extras.items():
        if key in field_map:
            field_kwargs[field_map[key]] = value
        else:
            schema_extras[key] = value

    root_annotation = base_cls.model_fields["root"].annotation
    if field_kwargs:
        root_annotation = Annotated[root_annotation, Field(**field_kwargs)]

    # Give the subclass a unique title so that multiple constrained variants of
    # the same base type (e.g. two FloatValue fields with different bounds in a
    # DataValue) produce distinct $defs keys instead of colliding.
    digest = json_digest(extras)
    unique_title = f"{base_cls.__name__}_{digest}"

    config_updates: dict[str, Any] = {"title": unique_title}
    if schema_extras:
        config_updates["json_schema_extra"] = schema_extras
    namespace: dict[str, Any] = {
        "__annotations__": {"root": root_annotation},
        "model_config": base_cls.model_config | ConfigDict(**config_updates),
    }

    return type(base_cls.__name__, (base_cls,), namespace, register=False)


def _build_constrained_sequence_cls(
    item_type: type,
    extras: dict[str, Any],
) -> type:
    """
    Return a constrained SequenceValue subclass that remains a proper Pydantic
    generic specialisation.  Unlike calling _build_constrained_cls on an already-
    concrete SequenceValue[T], this creates an intermediate generic class and then
    binds the item type, so __pydantic_generic_metadata__ has origin=<subclass>
    (not None).  That keeps get_origin_and_args() — and therefore the casting
    system — working correctly.
    """
    field_kwargs: dict[str, Any] = {}
    schema_extras: dict[str, Any] = {}
    for key, value in extras.items():
        if key in _SEQUENCE_FIELD_MAP:
            field_kwargs[_SEQUENCE_FIELD_MAP[key]] = value
        else:
            schema_extras[key] = value

    # Include item type name so that SequenceValue[FloatValue] and
    # SequenceValue[IntegerValue] with identical extras get distinct $defs keys.
    digest = json_digest({"item_type": item_type.__name__, **extras})
    unique_title = f"SequenceValue_{digest}"

    namespace: dict[str, Any] = {}
    if field_kwargs:
        namespace["__annotations__"] = {
            "root": Annotated[Sequence[_T_item], Field(**field_kwargs)]  # type: ignore[valid-type]
        }
    config_updates: dict[str, Any] = {"title": unique_title}
    if schema_extras:
        config_updates["json_schema_extra"] = schema_extras
    namespace["model_config"] = SequenceValue.model_config | ConfigDict(
        **config_updates
    )

    cls = type("SequenceValue", (SequenceValue[_T_item],), namespace, register=False)  # type: ignore[valid-type]
    return cls[item_type]  # type: ignore[index]


def _build_constrained_map_cls(
    item_type: type,
    extras: dict[str, Any],
) -> type:
    """
    Return a constrained StringMapValue subclass using the same generic approach
    as _build_constrained_sequence_cls.
    """
    field_kwargs: dict[str, Any] = {}
    schema_extras: dict[str, Any] = {}
    for key, value in extras.items():
        if key in _MAP_FIELD_MAP:
            field_kwargs[_MAP_FIELD_MAP[key]] = value
        else:
            schema_extras[key] = value

    digest = json_digest({"item_type": item_type.__name__, **extras})
    unique_title = f"StringMapValue_{digest}"

    namespace: dict[str, Any] = {}
    if field_kwargs:
        namespace["__annotations__"] = {
            "root": Annotated[Mapping[str, _T_item], Field(**field_kwargs)]  # type: ignore[valid-type]
        }
    config_updates: dict[str, Any] = {"title": unique_title}
    if schema_extras:
        config_updates["json_schema_extra"] = schema_extras
    namespace["model_config"] = StringMapValue.model_config | ConfigDict(
        **config_updates
    )

    cls = type("StringMapValue", (StringMapValue[_T_item],), namespace, register=False)  # type: ignore[valid-type]
    return cls[item_type]  # type: ignore[index]


class BaseValueSchema(ImmutableBaseModel):
    model_config: ClassVar[ConfigDict] = ConfigDict(
        extra="allow",
        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
    value_type: str | 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 Python attribute names from the
        object, whereas it is expecting their aliased forms.
        To fix this, we move the properties to the expected keys.
        """
        if not isinstance(data, Mapping):
            return data
        data = dict(data)
        if "$defs" not in data and "defs" in data:
            data["$defs"] = data.pop("defs")
        # Accept "x-value-type" (wire name) and map it to the Python field name.
        if "value_type" not in data and "x-value-type" in data:
            data["value_type"] = data.pop("x-value-type")
        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)

        # Rename value_type -> x-value-type for wire format before filtering.
        if "value_type" in data:
            data["x-value-type"] = data.pop("value_type")

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

        if hasattr(self, "__pydantic_fields_set__"):
            for key in self.__pydantic_fields_set__:
                # Translate the Python field name to the wire name.
                keys_to_keep.add("x-value-type" if key == "value_type" else key)

        # 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")

        # Keep x-value-type if it has a value, regardless of __pydantic_fields_set__.
        if data.get("x-value-type") is not None:
            keys_to_keep.add("x-value-type")

        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 x-value-type field matches a registered Value class and the
        schema carries no additional fields, that class is returned directly.
        This lets users provide { "x-value-type": "CustomValue" } as a
        shorthand for the entire CustomValue schema.

        If extra fields are present alongside x-value-type (e.g. constraints
        such as ``minimum`` or ``maxLength``), the schema is built normally so
        that a constrained subclass is produced and the constraints are not lost.

        References, if any, are resolved using self.defs first, then any
        extra_defs in order of decreasing precedence.
        """
        if not self.model_extra:
            value_cls = ValueRegistry.DEFAULT.load_value(self)
            if value_cls is not None:
                return value_cls
        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: Final[Literal["boolean"]]

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


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

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


class IntegerValueSchema(BaseValueSchema):
    type: Final[Literal["integer"]]

    @override
    def build_value_cls(
        self,
        *extra_defs: Mapping[str, ValueSchema],
    ) -> type[IntegerValue]:
        return _build_constrained_cls(
            IntegerValue, _NUMERIC_FIELD_MAP, self.model_extra or {}
        )


class FloatValueSchema(BaseValueSchema):
    type: Final[Literal["number"]]

    @override
    def build_value_cls(
        self,
        *extra_defs: Mapping[str, ValueSchema],
    ) -> type[FloatValue]:
        return _build_constrained_cls(
            FloatValue, _NUMERIC_FIELD_MAP, self.model_extra or {}
        )


class StringValueSchema(BaseValueSchema):
    type: Final[Literal["string"]]
    enum: Sequence[str] | None = None
    pattern: str | None = None

    @override
    def build_value_cls(
        self,
        *extra_defs: Mapping[str, ValueSchema],
    ) -> type[StringValue]:
        # `pattern` is a declared field so it doesn't appear in model_extra;
        # merge it in explicitly so the constraint helper can see it.
        extras = dict(self.model_extra or {})
        if self.pattern is not None:
            extras["pattern"] = self.pattern
        return _build_constrained_cls(StringValue, _STRING_FIELD_MAP, extras)


class SequenceValueSchema(BaseValueSchema):
    type: Final[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)
        extras = dict(self.model_extra or {})
        if not extras:
            return SequenceValue[T]
        return _build_constrained_sequence_cls(T, extras)


class StringMapValueSchema(BaseValueSchema):
    type: Final[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:
            item_type = Value
        else:
            item_type = self.additionalProperties.to_value_cls(self.defs, *extra_defs)
        extras = dict(self.model_extra or {})
        if not extras:
            return StringMapValue[item_type]
        return _build_constrained_map_cls(item_type, extras)


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

    type: Final[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


class FieldSchemaMappingValue(StringMapValue[ValueSchemaValue]):
    def to_data_schema(self, title: str) -> DataValueSchema:
        return DataValueSchema(
            type="object",
            title=title,
            properties={k: v.root for k, v in self.root.items()},
            additionalProperties=False,
        )

    @classmethod
    def from_fields(cls, **fields: ValueType) -> Self:
        return cls(
            {
                name: ValueSchemaValue(vtype.to_value_schema())
                for name, vtype in fields.items()
            }
        )


__all__ = [
    "BooleanValueSchema",
    "DataValueSchema",
    "FieldSchemaMappingValue",
    "FloatValueSchema",
    "IntegerValueSchema",
    "NullValueSchema",
    "ReferenceValueSchema",
    "SequenceValueSchema",
    "StringMapValueSchema",
    "StringValueSchema",
    "UnionValueSchema",
    "ValueSchema",
    "ValueSchemaValue",
]
