from dataclasses import dataclass, field, replace
from functools import cached_property
from typing import Any, Iterable, Literal, Protocol, Self, overload

import z3

from fest.symstructs import BoundsFormula, Fun, Rel, SymFormula, SymTerm
from fest.theories import TheoryName

from .template_free_variables import TemplateFreeScope

type Bound = Literal["bound"]


@dataclass(eq=True, frozen=True)
class Domain:
    regular: int
    summary: int

    @classmethod
    @overload
    def create(cls, regular: int, /) -> Self: ...

    @classmethod
    @overload
    def create(cls, regular: int, summary: int, /) -> Self: ...

    @classmethod
    @overload
    def create(cls, domain: tuple[int, int], /) -> Self: ...

    @classmethod
    @overload
    def create(cls, domain: Self, /) -> Self: ...

    @classmethod
    def create(
        cls, regular_or_tuple_or_domain: int | tuple[int, int] | Self, summary: int = 0
    ) -> Self:
        if isinstance(regular_or_tuple_or_domain, int):
            return cls(regular_or_tuple_or_domain, summary)
        elif isinstance(regular_or_tuple_or_domain, tuple):
            return cls(*regular_or_tuple_or_domain)
        else:
            return regular_or_tuple_or_domain


type DomainLike = Domain | int | tuple[int, int]
type SortLike = str | z3.SortRef


class _TemplateTermCallable(Protocol):
    def __call__(
        self, scope: TemplateFreeScope, /, *indices: z3.ExprRef
    ) -> z3.ExprRef: ...


@dataclass(eq=True, frozen=True)
class TemplateTerm:
    name: str
    arity: int
    callable: _TemplateTermCallable = field(compare=False)

    def __call__(self, scope: TemplateFreeScope, /, *indices: z3.ExprRef) -> z3.ExprRef:
        return self.callable(scope, *indices)


class _TemplateFormulaCallable(Protocol):
    def __call__(
        self, scope: TemplateFreeScope, /, *indices: z3.ExprRef
    ) -> z3.BoolRef: ...


@dataclass(eq=True, frozen=True)
class TemplateFormula:
    name: str
    arity: int
    callable: _TemplateFormulaCallable = field(compare=False)

    def __call__(self, scope: TemplateFreeScope, /, *indices: z3.ExprRef) -> z3.BoolRef:
        return self.callable(scope, *indices)


@dataclass(eq=True, frozen=True)
class Identifier:
    name: str
    arity: int

    def prefix(self, prefix: str) -> str:
        return prefix + self.name


@dataclass(kw_only=True)
class Template:
    domain: dict[z3.SortRef, Domain] = field(default_factory=dict)
    bounds: set[TemplateFormula] = field(default_factory=set)
    terms: set[TemplateTerm] = field(default_factory=set)
    formulas: set[TemplateFormula] = field(default_factory=set)
    theory_name: TheoryName = "lia"

    @cached_property
    def named_bounds(self) -> dict[str, TemplateFormula]:
        return {formula.name: formula for formula in self.bounds}

    @cached_property
    def named_terms(self) -> dict[Identifier, TemplateTerm]:
        return {Identifier(term.name, term.arity): term for term in self.terms}

    @cached_property
    def named_formulas(self) -> dict[Identifier, TemplateFormula]:
        return {
            Identifier(formula.name, formula.arity): formula
            for formula in self.formulas
        }

    @overload
    def for_domain(self, domain: dict[z3.SortRef, DomainLike], /) -> Self: ...

    @overload
    def for_domain(self, /, **kwargs: DomainLike) -> Self: ...

    def for_domain(
        self,
        domain: dict[z3.SortRef, DomainLike] | None = None,
        /,
        **kwargs: DomainLike,
    ) -> Self:
        source: Iterable[tuple[SortLike, DomainLike]]
        if domain is not None:
            source = domain.items()
        else:
            source = kwargs.items()

        return replace(
            self,
            domain={_get_sort(sort): Domain.create(domain) for sort, domain in source},
        )

    def __or__(self, other: Self) -> Self:
        return replace(
            self,
            domain=self.domain | other.domain,
            bounds=self.bounds | other.bounds,
            terms=self.terms | other.terms,
            formulas=self.formulas | other.formulas,
        )

    def apply_bounds(
        self, index: z3.ExprRef, scope: TemplateFreeScope
    ) -> dict[str, z3.BoolRef]:
        return {
            identifier: bound_formula(scope, index)
            for identifier, bound_formula in self.named_bounds.items()
        }

    def apply_terms(
        self, *indices: z3.ExprRef, scope: TemplateFreeScope
    ) -> dict[str, z3.ExprRef]:
        return {
            identifier.name: term(scope, *indices)
            for identifier, term in self.named_terms.items()
            if identifier.arity == len(indices)
        }

    def apply_formulas(
        self, *indices: z3.ExprRef, scope: TemplateFreeScope
    ) -> dict[str, z3.BoolRef]:
        return {
            identifier.name: formula(scope, *indices)
            for identifier, formula in self.named_formulas.items()
            if identifier.arity == len(indices)
        }

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

    @overload
    def __getitem__(self, item: tuple[Fun, str, int, TemplateFreeScope]) -> SymTerm: ...

    @overload
    def __getitem__(
        self, item: tuple[Rel, str, int, TemplateFreeScope]
    ) -> SymFormula: ...

    def __getitem__(self, item: Any) -> BoundsFormula | SymTerm | SymFormula:
        assert isinstance(item, tuple), f"Unsupported key {item}"

        head, *rest = item

        if head == "bound":
            return self._get_bound(*rest)
        elif head == "fun":
            return self._get_term(*rest)
        elif head == "rel":
            return self._get_formula(*rest)

        assert False, f"Unsupported key {item}"

    def _get_bound(self, identifier: str, scope: TemplateFreeScope) -> BoundsFormula:
        bound = self.named_bounds[identifier]
        return lambda index: bound(scope, index)

    def _get_term(
        self, identifier: str, arity: int, scope: TemplateFreeScope
    ) -> SymTerm:
        term = self.named_terms[Identifier(identifier, arity)]
        return lambda *indices: term(scope, *indices)

    def _get_formula(
        self, identifier: str, arity: int, scope: TemplateFreeScope
    ) -> SymFormula:
        formula = self.named_formulas[Identifier(identifier, arity)]
        return lambda *indices: formula(scope, *indices)


def _get_sort(sort: SortLike) -> z3.SortRef:
    if isinstance(sort, str):
        return z3.DeclareSort(sort)
    else:
        return sort
