import contextlib
import tomllib
from collections.abc import Callable, Iterator
from dataclasses import dataclass, field, replace
from functools import cache, cached_property
from importlib.resources import files
from inspect import signature
from pathlib import Path
from types import ModuleType
from typing import (
    IO,
    Any,
    ContextManager,
    Literal,
    Protocol,
    Self,
    TypeGuard,
    get_args,
    get_origin,
    overload,
)

from common import assets


class SupportsReadBytes(Protocol):
    def open(self, mode: Literal["rb"]) -> IO[bytes]: ...


@dataclass(frozen=True)
class Source:
    values: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_toml(cls, source: str | SupportsReadBytes | IO[bytes]) -> Self:
        if isinstance(source, str):
            return cls.flatten(tomllib.loads(source))
        elif isinstance(source, IO):
            return cls.flatten(tomllib.load(source))
        else:
            return cls.flatten(tomllib.load(source.open("rb")))

    @classmethod
    def flatten(cls, values: dict[str, Any], *, prefix: str | None = None) -> Self:
        if prefix is None:
            prefix = ""
        else:
            prefix = prefix + "."

        return cls(_flatten(values, prefix=prefix))

    @cache
    def __getitem__(self, item: str) -> Any:
        return self.values[item]

    @cache
    def __contains__(self, item: str) -> bool:
        return item in self.values

    def __iter__(self) -> Iterator[str]:
        return iter(self.values)

    def items(self) -> Iterator[tuple[str, Any]]:
        return iter(self.values.items())

    def __hash__(self) -> int:  # needed for cache
        return id(self)


type DumpedSource = dict[str, Any]
type DumpedRegistry = tuple[list[DumpedSource], list[DumpedSource]]


@dataclass(eq=True, unsafe_hash=True)
class Registry:
    user_sources: tuple[Source, ...]
    default_sources: tuple[Source, ...]

    @classmethod
    def load(cls, source: str | Path | IO[bytes] = Path("config.toml")) -> Self:
        try:
            if isinstance(source, str):
                loaded_source = Source.from_toml(source)
            elif isinstance(source, Path):
                if source.exists():
                    loaded_source = Source.from_toml(source)
                else:
                    loaded_source = Source()
            else:
                loaded_source = Source.from_toml(source)

            return cls((loaded_source,), (cls.default_source(),))
        except tomllib.TOMLDecodeError:
            assert False, f"Could not read config from {source}"

    @classmethod
    def default_source(cls, module: ModuleType = assets) -> Source:
        return Source.from_toml(files(module) / "default.config.toml")

    def dump(self) -> DumpedRegistry:
        return list(source.values for source in self.user_sources), list(
            source.values for source in self.default_sources
        )

    def load_from_dump(self, registry: DumpedRegistry) -> None:
        user_sources, default_sources = registry
        self.user_sources = tuple(Source(source) for source in user_sources)
        self.default_sources = tuple(Source(source) for source in default_sources)

    def add_default(self, source: Source | ModuleType | Path) -> None:
        if isinstance(source, Path):
            source = Source.from_toml(source)
        elif not isinstance(source, Source):
            source = self.default_source(source)

        self.default_sources = (source,) + self.default_sources

    def add_user(self, source: Source | ModuleType | Path) -> None:
        if isinstance(source, Path):
            source = Source.from_toml(source)
        elif not isinstance(source, Source):
            source = self.default_source(source)

        self.user_sources = (source,) + self.user_sources

    @overload
    def override(self, /, **kwargs: Any) -> ContextManager[None]: ...

    @overload
    def override(self, source: Source | dict[str, Any], /) -> ContextManager[None]: ...

    @contextlib.contextmanager
    def override(
        self, source: Source | dict[str, Any] | None = None, /, **kwargs: Any
    ) -> Iterator[None]:
        if source is None:
            source = Source(kwargs)
        elif not isinstance(source, Source):
            source = Source.flatten(source)

        user_sources = self.user_sources
        self.add_user(source)
        try:
            yield None
        finally:
            self.user_sources = user_sources

    @property
    def sources(self) -> tuple[Source, ...]:
        return (*self.user_sources, *self.default_sources)

    def __getitem__(self, item: str) -> Any:
        for source in self.sources:
            if item in source:
                return source[item]

        return None

    def __iter__(self) -> Iterator[str]:
        for source in self.sources:
            yield from source

    def items(self) -> Iterator[tuple[str, Any]]:
        for source in self.sources:
            yield from source.items()

    def __getattr__(self, item: str) -> "ConfigValue":
        return ConfigValue(item, self)

    def __call__[T, **P](
        self, **configs: "ConfigValue | TypedConfigValue[Any]"
    ) -> Callable[[Callable[P, T]], Callable[P, T]]:
        def wrapper(fun: Callable[P, T]) -> Callable[P, T]:
            sig = signature(fun)
            for kwarg, config in configs.items():
                assert (
                    kwarg in sig.parameters
                ), f"`{fun.__name__}` has no keyword argument `{kwarg}`"

                param = sig.parameters[kwarg]

                assert (
                    param.kind is param.POSITIONAL_OR_KEYWORD
                    or param.kind is param.KEYWORD_ONLY
                ), f"`{kwarg}` cannot be passed as a keyword argument"

                assert (
                    param.default is not param.empty
                ), f"`{kwarg}` has no default value"

            def wrapped(*args: P.args, **kwargs: P.kwargs) -> T:
                for kwarg, config in configs.items():
                    if kwarg not in kwargs:
                        if isinstance(config, TypedConfigValue):
                            kwargs[kwarg] = config()
                        else:
                            kwargs[kwarg] = config(sig.parameters[kwarg].annotation)

                return fun(*args, **kwargs)

            return wrapped

        return wrapper


@dataclass(eq=True, frozen=True)
class ConfigValue:
    _name: str
    _registry: Registry

    @cached_property
    def _value(self) -> Any:
        return self._registry[self._name]

    @overload
    def override(self, values: dict[str, Any], /) -> ContextManager[None]: ...

    @overload
    def override(self, value: Any, /) -> ContextManager[None]: ...

    @overload
    def override(self, /, **kwargs: Any) -> ContextManager[None]: ...

    def override(
        self, values: Any | None = None, /, **kwargs: Any
    ) -> ContextManager[None]:
        if isinstance(values, dict):
            return self._registry.override(Source.flatten(values, prefix=self._name))
        elif values is not None:
            return self._registry.override(Source.flatten({self._name: values}))
        else:
            return self._registry.override(Source.flatten(kwargs, prefix=self._name))

    def __bool__(self) -> bool:
        return self(bool)

    def __int__(self) -> int:
        return self(int)

    def __float__(self) -> float:
        return self(float)

    def __str__(self) -> str:
        return self(str)

    def __repr__(self) -> str:
        return f"ConfigValue <{self._name}>"

    def __getattr__(self, item: str) -> Self:
        return replace(self, _name=self._name + "." + item)

    def __getitem__[T](self, value_type: type[T]) -> "TypedConfigValue[T]":
        return TypedConfigValue(self, value_type)

    def __call__[T](self, value_type: type[T]) -> T:
        assert _check_type(
            self._value, value_type
        ), f"Expected {_type_name(value_type)} config for {self._name}, got {self._value}"

        return self._value


@dataclass(frozen=True, slots=True)
class TypedConfigValue[T]:
    value: ConfigValue
    value_type: type[T]

    def __str__(self) -> str:
        return f"(config {self.value._name} = {self()})"

    def __call__(self) -> T:
        return self.value(self.value_type)


def _type_name(type_to_name: type[Any]) -> str:
    if (origin := get_origin(type_to_name)) is not None:
        return (
            str(origin.__name__)
            + "["
            + ", ".join(_type_name(arg) for arg in get_args(type_to_name))
            + "]"
        )

    return type_to_name.__name__


def _check_type[T](value: Any, value_type: type[T]) -> TypeGuard[T]:
    origin = get_origin(value_type)

    if origin is None:
        return isinstance(value, value_type)

    if origin is list:
        args = get_args(value_type)
        if len(args) == 1 and all(_check_type(element, args[0]) for element in value):
            return True

    return False


def _flatten(values: dict[str, Any], prefix: str) -> dict[str, Any]:
    flattened: dict[str, Any] = {}

    def _flatten_rec(sub_values: dict[str, Any], current_prefix: str) -> None:
        nonlocal flattened
        for key, value in sub_values.items():
            full_key = current_prefix + key
            if isinstance(value, dict):
                _flatten_rec(value, full_key + ".")
            else:
                flattened[full_key] = value

    _flatten_rec(values, prefix)
    return flattened
