#!/usr/bin/env python3

# Libervia models custom types
# Copyright (C) 2009-2024 Jérôme Poisson (goffi@goffi.org)

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from typing import TYPE_CHECKING, Annotated, Any, TypeAlias, cast

from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler
from pydantic.json_schema import JsonSchemaValue
from pydantic_core import core_schema
from twisted.words.protocols.jabber.jid import JID
from twisted.words.xish.domish import Element
from libervia.backend.tools.xml_tools import parse


class _DomishElementType:
    """Use Twisted's domish.Element in Python type, and serialize it to str

    In Python, both Element and str are accepted, str are converted to Element.
    Serialization is done using toXml() method, and parsing using xml_tools.parse.
    """

    @staticmethod
    def validate_element(value: str) -> Element:
        return parse(value)

    @staticmethod
    def serialize_element(element: Element) -> str:
        return element.toXml()

    @classmethod
    def __get_pydantic_core_schema__(
        cls, source_type: Any, handler: GetCoreSchemaHandler
    ) -> core_schema.CoreSchema:
        return core_schema.json_or_python_schema(
            json_schema=core_schema.no_info_after_validator_function(
                cls.validate_element, core_schema.str_schema()
            ),
            python_schema=core_schema.union_schema(
                [
                    core_schema.is_instance_schema(Element),
                    core_schema.no_info_after_validator_function(
                        cls.validate_element, core_schema.str_schema()
                    ),
                ]
            ),
            serialization=core_schema.plain_serializer_function_ser_schema(
                cls.serialize_element, when_used="json"
            ),
        )

    @classmethod
    def __get_pydantic_json_schema__(
        cls, schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
    ) -> JsonSchemaValue:
        json_schema = handler(schema)
        json_schema.update(
            {
                "format": "xml",
                "description": "A valid XML element that can be parsed into a "
                "domish.Element",
            }
        )
        return json_schema


class _StrictDomishElementType(_DomishElementType):
    """Strict version of DomishElementType which only accept Element in Python."""

    @classmethod
    def __get_pydantic_core_schema__(
        cls, source_type: Any, handler: GetCoreSchemaHandler
    ) -> core_schema.CoreSchema:
        return core_schema.json_or_python_schema(
            json_schema=core_schema.no_info_after_validator_function(
                cls.validate_element, core_schema.str_schema()
            ),
            python_schema=core_schema.is_instance_schema(Element),
            serialization=core_schema.plain_serializer_function_ser_schema(
                cls.serialize_element, when_used="json"
            ),
        )


class _JIDType:
    """Use Twisted's JID in Python type, and serialize it to str

    In Python, both JID and str are accepted, str are converted to JID.
    """

    @staticmethod
    def validate_jid(value: str) -> JID:
        return JID(value)

    @staticmethod
    def serialize_jid(jid: JID) -> str:
        return str(jid)

    @classmethod
    def __get_pydantic_core_schema__(
        cls, source_type: Any, handler: GetCoreSchemaHandler
    ) -> core_schema.CoreSchema:
        return core_schema.json_or_python_schema(
            json_schema=core_schema.no_info_after_validator_function(
                cls.validate_jid, core_schema.str_schema()
            ),
            python_schema=core_schema.union_schema(
                [
                    core_schema.is_instance_schema(JID),
                    core_schema.no_info_after_validator_function(
                        cls.validate_jid, core_schema.str_schema()
                    ),
                ]
            ),
            serialization=core_schema.plain_serializer_function_ser_schema(
                cls.serialize_jid, when_used="json"
            ),
        )

    @classmethod
    def __get_pydantic_json_schema__(
        cls, schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
    ) -> JsonSchemaValue:
        json_schema = handler(schema)
        json_schema.update(
            {
                "format": "jid",
                "description": "A valid Jabber ID (JID) in the format "
                '"user@domain/resource"',
            }
        )
        return json_schema


class _StrictJIDType(_JIDType):
    """Strict version of JIDType which only accept JID in Python."""

    @classmethod
    def __get_pydantic_core_schema__(
        cls, source_type: Any, handler: GetCoreSchemaHandler
    ) -> core_schema.CoreSchema:
        return core_schema.json_or_python_schema(
            json_schema=core_schema.no_info_after_validator_function(
                cls.validate_jid, core_schema.str_schema()
            ),
            python_schema=core_schema.is_instance_schema(JID),
            serialization=core_schema.plain_serializer_function_ser_schema(
                cls.serialize_jid, when_used="json"
            ),
        )


# We annotate the types so static type checker understand them as JID.
JIDType = Annotated[JID, _JIDType]
StrictJIDType = Annotated[JID, _StrictJIDType]
DomishElementType = Annotated[Element, _DomishElementType]
StrictDomishElementType = Annotated[Element, _StrictDomishElementType]