from base64 import b64decode, b64encode
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple, Type
from uuid import UUID

from jinja2 import Markup

from quart.json import dumps, loads


class JSONTag:
    key: Optional[str] = None

    def __init__(self, serializer: 'TaggedJSONSerializer') -> None:
        self.serializer = serializer

    def check(self, value: Any) -> bool:
        raise NotImplementedError

    def to_json(self, value: Any) -> Any:
        raise NotImplementedError

    def to_python(self, value: str) -> Any:
        raise NotImplementedError

    def tag(self, value: Any) -> Any:
        return {self.key: self.to_json(value)}


class TagDict(JSONTag):
    key = ' di'

    def check(self, value: Any) -> bool:
        return (
            isinstance(value, dict) and len(value) == 1 and
            next(iter(value)) in self.serializer.tags
        )

    def to_json(self, value: Any) -> Dict[str, Any]:
        key = next(iter(value))
        return {key + '__': self.serializer.tag(value[key])}

    def to_python(self, value: str) -> Dict[str, Any]:
        key, item = next(iter(value))
        return {key[:-2]: item}


class PassDict(JSONTag):

    def check(self, value: Any) -> bool:
        return isinstance(value, dict)

    def to_json(self, value: Any) -> Dict[str, Any]:
        return {key: self.serializer.tag(item) for key, item in value.items()}

    tag = to_json


class TagTuple(JSONTag):
    key = ' t'

    def check(self, value: Any) -> bool:
        return isinstance(value, tuple)

    def to_json(self, value: Tuple[Any]) -> List[Any]:
        return [self.serializer.tag(item) for item in value]

    def to_python(self, value: Any) -> Tuple[Any, ...]:
        return tuple(value)


class PassList(JSONTag):

    def check(self, value: Any) -> bool:
        return isinstance(value, list)

    def to_json(self, value: List[Any]) -> List[Any]:
        return [self.serializer.tag(item) for item in value]

    tag = to_json


class TagBytes(JSONTag):
    key = ' b'

    def check(self, value: Any) -> bool:
        return isinstance(value, bytes)

    def to_json(self, value: bytes) -> str:
        return b64encode(value).decode('ascii')

    def to_python(self, value: str) -> bytes:
        return b64decode(value)


class TagMarkup(JSONTag):
    key = ' m'

    def check(self, value: Any) -> bool:
        return callable(getattr(value, '__html__', None))

    def to_json(self, value: Any) -> str:
        return str(value.__html__())

    def to_python(self, value: str) -> Markup:
        return Markup(value)


class TagUUID(JSONTag):
    key = ' u'

    def check(self, value: Any) -> bool:
        return isinstance(value, UUID)

    def to_json(self, value: Any) -> str:
        return value.hex

    def to_python(self, value: str) -> UUID:
        return UUID(value)


def _parse_datetime(value: str) -> datetime:
    try:
        return datetime.strptime(value, "%Y-%m-%dT%H:%M:%S.%f")
    except ValueError:
        return datetime.strptime(value, "%Y-%m-%dT%H:%M:%S.%f%z")


class TagDateTime(JSONTag):
    key = ' d'

    def check(self, value: Any) -> bool:
        return isinstance(value, datetime)

    def to_json(self, value: datetime) -> str:
        return value.isoformat(timespec='microseconds')

    def to_python(self, value: str) -> datetime:
        return _parse_datetime(value)


class TaggedJSONSerializer:

    default_tags = [
        TagDict, PassDict, TagTuple, PassList, TagBytes, TagMarkup, TagUUID,
        TagDateTime,
    ]

    def __init__(self) -> None:
        self.tags: Dict[str, JSONTag] = {}
        self.order: List[JSONTag] = []

        for tag_class in self.default_tags:
            self.register(tag_class)

    def register(
            self, tag_class: Type[JSONTag], force: bool=False, index: Optional[int]=None,
    ) -> None:
        tag = tag_class(self)
        key = tag.key

        if key is not None:
            if not force and key in self.tags:
                raise KeyError(f"Tag '{key}' is already registered.")

            self.tags[key] = tag

        if index is None:
            self.order.append(tag)
        else:
            self.order.insert(index, tag)

    def tag(self, value: Any) -> Dict[str, Any]:
        for tag in self.order:
            if tag.check(value):
                return tag.tag(value)

        return value

    def untag(self, value: Dict[str, Any]) -> Any:
        if len(value) != 1:
            return value

        key = next(iter(value))

        if key not in self.tags:
            return value

        return self.tags[key].to_python(value[key])

    def dumps(self, value: Any) -> str:
        return dumps(self.tag(value), separators=(',', ':'))

    def loads(self, value: str) -> Any:
        return loads(value, object_hook=self.untag)
