from abc import ABC, abstractmethod
from collections.abc import Iterator
from dataclasses import dataclass, field

import z3

from fest.theories import BaseTheory, TheoryName


class TemplateFreeVariables(ABC):
    @abstractmethod
    def __getitem__(self, item: tuple[str, str]) -> z3.ExprRef: ...

    @property
    @abstractmethod
    def theory_name(self) -> TheoryName: ...

    def scope(self, scope: str) -> "TemplateFreeScope":
        return TemplateFreeScope(scope, self)


@dataclass
class TemplateFreeVariablesGenerator(TemplateFreeVariables):
    variables: set[str] = field(default_factory=set, init=False)
    theory: BaseTheory

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

    def __getitem__(self, variable: tuple[str, str]) -> z3.ExprRef:
        prefix, suffix = variable
        variable_name = prefix + suffix
        const = self.theory.to_base_const(variable_name)
        self.variables.add(variable_name)
        return const

    @property
    def theory_name(self) -> TheoryName:
        return self.theory.name()


@dataclass
class TemplateFreeScope:
    prefix: str
    generator: TemplateFreeVariables

    def __getitem__(self, item: str) -> z3.ExprRef:
        return self.generator[self.prefix, item]

    @property
    def theory_name(self) -> TheoryName:
        return self.generator.theory_name
