import json
from typing import Any, List, Optional, Type, Union

from django.core.exceptions import ValidationError
from django.db.models import JSONField
from nwon_baseline.pydantic import pydantic_model_to_dict
from nwon_baseline.typings import AnyDict
from pydantic import BaseModel
from pydantic import ValidationError as PydanticValidationError


class PydanticJsonField(JSONField):
    """
    An extension of a JSONField that uses pydantic models for data validation.

    If data is valid for multiple given models the first valid model is returned.
    """

    def __init__(
        self, *args, pydantic_models: Optional[List[Type[BaseModel]]] = None, **kwargs
    ):
        self.pydantic_models: List[Type[BaseModel]] = pydantic_models or []
        kwargs.pop("pydantic_models", None)
        super().__init__(*args, **kwargs)

    def _validate_schema(
        self, value: Union[str, AnyDict, BaseModel, None]
    ) -> BaseModel:
        """
        Validates the value against the provided Pydantic models.
        Returns the first valid Pydantic model instance or raises ValidationError.
        """

        # Skip validation during fake migrations
        if self.model.__module__ == "__fake__":
            return value  # type: ignore

        errors: List[str] = []

        for model in self.pydantic_models:
            try:
                if isinstance(value, str):
                    return model.model_validate_json(value)
                else:
                    return model.model_validate(value)

            except PydanticValidationError as exc:
                errors.append(
                    f"JSON does not fit Pydantic model {model.__name__} {format(exc)}"
                )

        # Raise a validation error if no model matches
        raise ValidationError(
            f"Value does not match any Pydantic model: {errors}", code="invalid"
        )

    def validate(self, value: Optional[BaseModel], model_instance):
        if value is None:
            super().validate(value, model_instance)
        else:
            super().validate(pydantic_model_to_dict(value), model_instance)

        self._validate_schema(value)

    def pre_save(self, model_instance, add):
        """
        Validates and serializes the value before saving it to the database.
        """
        value = super().pre_save(model_instance, add)

        if value is None:
            return value

        validated_model = self._validate_schema(value)

        # Serialize the Pydantic model back to JSON
        return pydantic_model_to_dict(validated_model)

    def from_db_value(
        self, value: Any, expression: Any, connection: Any
    ) -> Union[BaseModel, None]:
        """
        Converts the database value into a Pydantic model instance when retrieving it.
        """
        if value is None:
            return value

        # Validate and return the parsed Pydantic model
        return self._validate_schema(value)

    def to_python(self, value: Any) -> Union[BaseModel, None]:
        """
        Converts the JSON value to a Pydantic model instance in Python.
        """
        if value is None:
            return value

        return self._validate_schema(value)
