import operator
from collections.abc import Callable
from dataclasses import dataclass, field
from functools import reduce
from importlib import import_module
from pkgutil import walk_packages
from typing import Any, overload

from problems.problem import Problem, ProblemCallable

type ProblemDecorator = Callable[[ProblemCallable], ProblemCallable]

type MarkerValue = bool | int | str


@dataclass
class Registry:
    problems: dict[str, tuple[Problem, dict[str, MarkerValue]]] = field(
        default_factory=dict, init=False
    )
    loaded: bool = field(default=False, init=False)

    @overload
    def __call__(self, f: None, /) -> ProblemDecorator: ...
    @overload
    def __call__(self, f: ProblemCallable, /) -> ProblemCallable: ...

    @overload
    def __call__(self, **markers: MarkerValue) -> ProblemDecorator: ...

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        def wrap(f: ProblemCallable) -> ProblemCallable:
            name = f.__name__
            assert (
                name not in self.problems
            ), f"Two problems with the same name: `{name}`"
            self.problems[name] = (Problem(name, f), kwargs)
            return f

        if not args:
            return wrap
        else:
            return wrap(args[0])

    def __contains__(self, item: str) -> bool:
        if not self.loaded:
            self.load()

        return item in self.problems

    def __getitem__(self, item: str) -> Problem:
        if not self.loaded:
            self.load()

        return self.problems[item][0]

    def __len__(self) -> int:
        if not self.loaded:
            self.load()

        return len(self.problems)

    def load(self) -> None:
        self.loaded = True

        package = import_module(__package__)
        package_path = package.__path__
        package_name = package.__name__

        for _loader, module_name, _is_pkg in walk_packages(
            package_path, package_name + "."
        ):
            import_module(module_name)

    def get(self, **markers: MarkerValue) -> set[Problem]:
        if not self.loaded:
            self.load()

        return {
            f
            for f, f_markers in self.problems.values()
            if self._match(markers, f_markers)
        }

    @staticmethod
    def _match(
        markers1: dict[str, MarkerValue], markers2: dict[str, MarkerValue]
    ) -> bool:
        for key, value in markers1.items():
            if isinstance(value, bool):
                if not value and key in markers2 and markers2[key]:
                    return False
                if value and key in markers2 and not markers2[key]:
                    return False
                if value and key not in markers2:
                    return False
            else:
                if key not in markers2 or markers2[key] != value:
                    return False

        return True

    def markers(self) -> set[str]:
        if not self.loaded:
            self.load()

        return reduce(
            operator.or_,
            (set(markers.keys()) for _problem, markers in self.problems.values()),
        )
