from collections import defaultdict
from collections.abc import Iterator
from dataclasses import dataclass, field, replace
from functools import cached_property
from typing import Mapping, Self, overload

import z3

from common.simplifier import Simplifier


@dataclass(eq=True, frozen=True, order=True, slots=True)
class Guard:
    scope: str
    name: str

    def __str__(self) -> str:
        return f"{self.scope}:{self.name}"

    def to_bool(self) -> z3.BoolRef:
        return z3.Bool(str(self))


@dataclass
class Generator(Mapping[str, set[Guard]]):
    guards: dict[str, set[Guard]] = field(
        default_factory=lambda: defaultdict(set), init=False
    )

    @overload
    def __getitem__(self, item: str) -> set[Guard]: ...

    @overload
    def __getitem__(self, item: tuple[str, str]) -> Guard: ...

    def __getitem__(self, item: str | tuple[str, str]) -> set[Guard] | Guard:
        if isinstance(item, str):
            return self.guards[item]

        scope, name = item
        guard = Guard(scope, name)
        self.guards[scope].add(guard)
        return guard

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

    def __len__(self) -> int:
        return len(self.guards)


@dataclass(eq=True, frozen=True)
class Clause:
    guards: frozenset[Guard] = field(default_factory=frozenset)

    @classmethod
    def create(cls, guard: Guard) -> Self:
        return cls(frozenset({guard}))

    @cached_property
    def scopes(self) -> frozenset[str]:
        return frozenset({guard.scope for guard in self.guards})

    @cached_property
    def ordered_guards(self) -> tuple[Guard, ...]:
        return tuple(sorted(list(self.guards)))

    def add(self, guard: Guard) -> Self | None:
        if guard in self.guards:
            return self
        elif guard.scope in self.scopes:
            return None
        return replace(self, guards=self.guards | {guard})

    def to_bool(self, simplifier: Simplifier) -> z3.BoolRef:
        return simplifier.simplified_and(
            *(guard.to_bool() for guard in self.ordered_guards)
        )

    def __rand__(self, other: Self | None) -> Self | None:
        return self & other

    def __and__(self, other: Self | None) -> Self | None:
        if other is None:
            return None

        guards = {}
        for guard in self.guards:
            guards[guard.scope] = guard

        for guard in other.guards:
            if guard.scope in guards and guards[guard.scope] != guard:
                return None

        return replace(self, guards=self.guards | other.guards)

    def __str__(self) -> str:
        return "{" + ", ".join(str(guard) for guard in self.ordered_guards) + "}"
