#!/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


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]