"""
The MIT License (MIT)

Copyright (c) 2021-present Disnake Development

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
DEALINGS IN THE SOFTWARE.
"""

from __future__ import annotations

import logging
import os
import warnings
from abc import ABC, abstractmethod
from collections import defaultdict
from pathlib import Path
from typing import (
    TYPE_CHECKING,
    Any,
    DefaultDict,
    Dict,
    Generic,
    Literal,
    Optional,
    Set,
    TypeVar,
    Union,
    overload,
)

from . import utils
from .custom_warnings import LocalizationWarning
from .enums import Locale
from .errors import LocalizationKeyError

if TYPE_CHECKING:
    from typing_extensions import Self

    LocalizedRequired = Union[str, "Localized[str]"]
    LocalizedOptional = Union[Optional[str], "Localized[Optional[str]]"]


__all__ = (
    "Localized",
    "Localised",
    "LocalizationValue",
    "LocalizationProtocol",
    "LocalizationStore",
)

MISSING = utils.MISSING

_log = logging.getLogger(__name__)


LocalizationsDict = Union[Dict[Locale, str], Dict[str, str]]
Localizations = Union[str, LocalizationsDict]

StringT = TypeVar("StringT", str, Optional[str], covariant=True)


# This is generic over `string`, as some localized strings can be optional, e.g. option descriptions.
# The basic idea for parameters is this:
#     abc: LocalizedRequired
#     xyz: LocalizedOptional = None
#
# With that, one may use `abc="somename"` and `abc=Localized("somename", ...)`,
# but not `abc=Localized(None, ...)`. All three work fine for `xyz` though.


class Localized(Generic[StringT]):
    """
    A container type used for localized parameters.

    Exactly one of ``key`` or ``data`` must be provided.

    There is an alias for this called ``Localised``.

    .. versionadded:: 2.5

    Parameters
    ----------
    string: Optional[:class:`str`]
        The default (non-localized) value of the string.
        Whether this is optional or not depends on the localized parameter type.
    key: :class:`str`
        A localization key used for lookups.
        Incompatible with ``data``.
    data: Union[Dict[:class:`.Locale`, :class:`str`], Dict[:class:`str`, :class:`str`]]
        A mapping of locales to localized values.
        Incompatible with ``key``.
    """

    __slots__ = ("string", "localizations")

    @overload
    def __init__(self: Localized[StringT], string: StringT, *, key: str):
        ...

    @overload
    def __init__(self: Localized[Optional[str]], *, key: str):
        ...

    @overload
    def __init__(
        self: Localized[StringT],
        string: StringT,
        *,
        data: Union[Optional[LocalizationsDict], LocalizationValue],
    ):
        ...

    @overload
    def __init__(
        self: Localized[Optional[str]],
        *,
        data: Union[Optional[LocalizationsDict], LocalizationValue],
    ):
        ...

    # note: `data` accepting `LocalizationValue` is intentionally undocumented,
    # as it's only meant to be used internally
    def __init__(
        self,
        string: StringT = None,
        *,
        key: str = MISSING,
        data: Union[Optional[LocalizationsDict], LocalizationValue] = MISSING,
    ):
        self.string: StringT = string

        if not (key is MISSING) ^ (data is MISSING):
            raise TypeError("Exactly one of `key` or `data` must be provided")
        if isinstance(data, LocalizationValue):
            self.localizations = data
        else:
            self.localizations = LocalizationValue(key if key is not MISSING else data)

    @overload
    @classmethod
    def _cast(cls, string: LocalizedOptional, required: Literal[False]) -> Localized[Optional[str]]:
        ...

    @overload
    @classmethod
    def _cast(cls, string: LocalizedRequired, required: Literal[True]) -> Localized[str]:
        ...

    @classmethod
    def _cast(cls, string: Union[Optional[str], Localized[Any]], required: bool) -> Localized[Any]:
        if not isinstance(string, Localized):
            string = cls(string, data=None)

        # enforce the `StringT` type at runtime
        if required and string.string is None:
            raise ValueError("`string` parameter must be provided")
        return string

    @overload
    def _upgrade(self, *, key: Optional[str]) -> Self:
        ...

    @overload
    def _upgrade(self, string: str, *, key: Optional[str] = None) -> Localized[str]:
        ...

    def _upgrade(
        self, string: Optional[str] = None, *, key: Optional[str] = None
    ) -> Localized[Any]:
        # update key if provided and not already set
        self.localizations._upgrade(key)

        # Only overwrite if not already set (`Localized()` parameter value takes precedence over function names etc.)
        # Note: not checking whether `string` is an empty string, to keep generic typing correct
        if not self.string and string is not None:
            self.string = string

        # this is safe, see above
        return self


Localised = Localized


class LocalizationValue:
    """
    Container type for (pending) localization data.

    .. versionadded:: 2.5
    """

    __slots__ = ("_key", "_data")

    def __init__(self, localizations: Optional[Localizations]):
        self._key: Optional[str]
        self._data: Optional[Dict[str, str]]

        if localizations is None:
            # no localization
            self._key = None
            self._data = None
        elif isinstance(localizations, str):
            # got localization key
            self._key = localizations
            self._data = MISSING  # not localized yet
        elif isinstance(localizations, dict):
            # got localization data
            self._key = None
            self._data = {str(k): v for k, v in localizations.items()}
        else:
            raise TypeError(f"Invalid localizations type: {type(localizations).__name__}")

    def _upgrade(self, key: Optional[str]) -> None:
        if not key:
            return

        # if empty, use new key
        if self._key is None and self._data is None:
            self._key = key
            self._data = MISSING
            return

        # if key is the same, ignore
        if self._key == key:
            return

        # at this point, the keys don't match, which either means that they're different strings,
        # or that there is no existing `_key` but `_data` is set
        raise ValueError("Can't specify multiple localization keys or dicts")

    def _link(self, store: LocalizationProtocol) -> None:
        """Loads localizations from the specified store if this object has a key."""
        if self._key is not None:
            self._data = store.get(self._key)

    @property
    def data(self) -> Optional[Dict[str, str]]:
        """Optional[Dict[:class:`str`, :class:`str`]]: A dict with a locale -> localization mapping, if available."""
        if self._data is MISSING:
            warnings.warn(
                "value was never localized, this is likely a library bug",
                LocalizationWarning,
                stacklevel=2,
            )
            return None
        return self._data

    def __eq__(self, other) -> bool:
        d1 = self.data
        d2 = other.data
        # consider values equal if they're both falsy, or actually equal
        # (it doesn't matter if localizations are `None` or `{}`)
        return (not d1 and not d2) or d1 == d2


class LocalizationProtocol(ABC):
    """
    Manages a key-value mapping of localizations.

    This is an abstract class, a concrete implementation is provided as :class:`LocalizationStore`.

    .. versionadded:: 2.5
    """

    @abstractmethod
    def get(self, key: str) -> Optional[Dict[str, str]]:
        """
        Returns localizations for the specified key.

        Parameters
        ----------
        key: :class:`str`
            The lookup key.

        Raises
        ------
        LocalizationKeyError
            May be raised if no localizations for the provided key were found,
            depending on the implementation.

        Returns
        -------
        Optional[Dict[:class:`str`, :class:`str`]]
            The localizations for the provided key.
            May return ``None`` if no localizations could be found.
        """
        raise NotImplementedError

    # subtypes don't have to implement this
    def load(self, path: Union[str, os.PathLike]) -> None:
        """
        Adds localizations from the provided path.

        Parameters
        ----------
        path: Union[:class:`str`, :class:`os.PathLike`]
            The path to the file/directory to load.

        Raises
        ------
        RuntimeError
            The provided path is invalid or couldn't be loaded
        """
        raise NotImplementedError

    # subtypes don't have to implement this
    def reload(self) -> None:
        """
        Clears localizations and reloads all previously loaded sources again.
        If an exception occurs, the previous data gets restored and the exception is re-raised.
        """
        pass


class LocalizationStore(LocalizationProtocol):
    """
    Manages a key-value mapping of localizations using ``.json`` files.

    .. versionadded:: 2.5

    Attributes
    ------------
    strict: :class:`bool`
        Specifies whether :meth:`.get` raises an exception if localizations for a provided key couldn't be found.
    """

    def __init__(self, *, strict: bool):
        self.strict = strict

        self._loc: DefaultDict[str, Dict[str, str]] = defaultdict(dict)
        self._paths: Set[Path] = set()

    def get(self, key: str) -> Optional[Dict[str, str]]:
        """
        Returns localizations for the specified key.

        Parameters
        ----------
        key: :class:`str`
            The lookup key.

        Raises
        ------
        LocalizationKeyError
            No localizations for the provided key were found.
            Raised only if :attr:`strict` is enabled, returns ``None`` otherwise.

        Returns
        -------
        Optional[Dict[:class:`str`, :class:`str`]]
            The localizations for the provided key.
            Returns ``None`` if no localizations could be found and :attr:`strict` is disabled.
        """

        data = self._loc.get(key)
        if data is None and self.strict:
            raise LocalizationKeyError(key)
        return data

    def load(self, path: Union[str, os.PathLike]) -> None:
        """
        Adds localizations from the provided path to the store.
        If the path points to a file, the file gets loaded.
        If it's a directory, all ``.json`` files in that directory get loaded (non-recursive).

        Parameters
        ----------
        path: Union[:class:`str`, :class:`os.PathLike`]
            The path to the file/directory to load.

        Raises
        ------
        RuntimeError
            The provided path is invalid or couldn't be loaded
        """

        path = Path(path)

        if path.is_file():
            self._load_file(path)
        elif path.is_dir():
            for file in path.glob("*.json"):
                if not file.is_file():
                    continue
                self._load_file(file)
        else:
            raise RuntimeError(f"Path '{path}' does not exist or is not a directory/file")

        self._paths.add(path)

    def reload(self) -> None:
        """
        Clears localizations and reloads all previously loaded files/directories again.
        If an exception occurs, the previous data gets restored and the exception is re-raised.
        See :func:`~LocalizationStore.load` for possible raised exceptions.
        """

        old = self._loc
        try:
            self._loc = defaultdict(dict)
            for path in self._paths:
                self.load(path)
        except Exception:
            # restore in case of error
            self._loc = old
            raise

    def _load_file(self, path: Path) -> None:
        try:
            if path.suffix != ".json":
                raise ValueError(f"not a .json file")
            locale = path.stem

            if not (api_locale := utils.as_valid_locale(locale)):
                raise ValueError(f"invalid locale '{locale}'")
            locale = api_locale

            data = utils._from_json(path.read_text("utf-8"))
            self._load_dict(data, locale)
            _log.debug(f"Loaded localizations from '{path}'")
        except Exception as e:
            raise RuntimeError(f"Unable to load '{path}': {e}") from e

    def _load_dict(self, data: Dict[str, str], locale: str) -> None:
        if not isinstance(data, dict) or not all(
            o is None or isinstance(o, str) for o in data.values()
        ):
            raise TypeError("data must be a flat dict with string/null values")
        for key, value in data.items():
            d = self._loc[key]  # always create dict, regardless of value
            if value is not None:
                d[locale] = value