from collections.abc import Callable

import z3

from fest.templates import TemplateFormula, TemplateTerm
from fest.templates.template_free_variables import TemplateFreeScope


def const_term(name: str, term: Callable[[], z3.ExprRef], arity: int) -> TemplateTerm:
    def wrapper(_scope: TemplateFreeScope, *indices: z3.ExprRef) -> z3.ExprRef:
        assert (
            len(indices) == arity
        ), f"Mismatch arity for term {name}: expected {arity} got {len(indices)}"
        return term()

    return TemplateTerm(name, arity, wrapper)


def const_formula(
    name: str, formula: Callable[[], z3.BoolRef], arity: int
) -> TemplateFormula:
    def wrapper(_scope: TemplateFreeScope, *indices: z3.ExprRef) -> z3.BoolRef:
        assert (
            len(indices) == arity
        ), f"Mismatch arity for formula {name}: expected {arity} got {len(indices)}"
        return formula()

    return TemplateFormula(name, arity, wrapper)


_TRUE = lambda: z3.BoolVal(True)
TRUE0 = const_formula("true", _TRUE, arity=0)
TRUE1 = const_formula("true", _TRUE, arity=1)
TRUE2 = const_formula("true", _TRUE, arity=2)
TRUE3 = const_formula("true", _TRUE, arity=3)

_FALSE = lambda: z3.BoolVal(False)
FALSE0 = const_formula("false", _FALSE, arity=0)
FALSE1 = const_formula("false", _FALSE, arity=1)
FALSE2 = const_formula("false", _FALSE, arity=2)
FALSE3 = const_formula("false", _FALSE, arity=3)
