import enum
from inspect import signature
from typing import Callable, overload

import z3

from ..symstructs import SymFormula, SymTerm
from . import Identifier, Template
from .template_free_variables import TemplateFreeScope

type _Term0 = Callable[[], z3.ArithRef]
type _Term1 = Callable[[z3.ArithRef], z3.ArithRef]
type _Term2 = Callable[[z3.ArithRef, z3.ArithRef], z3.ArithRef]
type _AnyTerm = _Term0 | _Term1 | _Term2


type _Formula0 = Callable[[], z3.BoolRef]
type _Formula1 = Callable[[z3.ArithRef], z3.BoolRef]
type _Formula2 = Callable[[z3.ArithRef, z3.ArithRef], z3.BoolRef]
type _Formula3 = Callable[[z3.ArithRef, z3.ArithRef, z3.ArithRef], z3.BoolRef]
type _AnyFormula = _Formula0 | _Formula1 | _Formula2 | _Formula3


type _ScopedFormula1 = Callable[[TemplateFreeScope, z3.ArithRef], z3.BoolRef]
type _ScopedFormula2 = Callable[
    [TemplateFreeScope, z3.ArithRef, z3.ArithRef], z3.BoolRef
]
type _ScopedFormula3 = Callable[
    [TemplateFreeScope, z3.ArithRef, z3.ArithRef, z3.ArithRef], z3.BoolRef
]
type _AnyScopedFormula = _ScopedFormula1 | _ScopedFormula2 | _ScopedFormula3


def bound(name: str, fun: _Formula1) -> Template:
    def wrapper(_scope: TemplateFreeScope, *indices: z3.ArithRef) -> z3.BoolRef:
        assert (
            len(indices) == 1
        ), f"Mismatch arity for bound {name}: expected 1 got {len(indices)}"
        return fun(*indices)

    return Template(bounds={name: wrapper})


def scoped_bound(name: str, fun: _ScopedFormula1) -> Template:
    def wrapper(scope: TemplateFreeScope, *indices: z3.ArithRef) -> z3.BoolRef:
        assert (
            len(indices) == 1
        ), f"Mismatch arity for bound {name}: expected 1 got {len(indices)}"
        return fun(scope, *indices)

    return Template(bounds={name: wrapper})


@overload
def term(name: str, fun: SymTerm, *, arity: int) -> Template: ...


@overload
def term(name: str, fun: _AnyTerm) -> Template: ...


def term(name: str, fun: _AnyTerm, *, arity: int | None = None) -> Template:
    if arity is None:
        assert callable(fun), "Must pass callable term when no arity is given"
        arity = len(signature(fun).parameters)

    def wrapper(_scope: TemplateFreeScope, *indices: z3.ArithRef) -> z3.ArithRef:
        assert (
            len(indices) == arity
        ), f"Mismatch arity for term {name}: expected {arity} got {len(indices)}"
        return fun(*indices)

    return Template(terms={Identifier(name, arity): wrapper})


@overload
def formula(name: str, fun: SymFormula, *, arity: int) -> Template: ...


@overload
def formula(name: str, fun: _AnyFormula) -> Template: ...


def formula(name: str, fun: _AnyFormula, *, arity: int | None = None) -> Template:
    if arity is None:
        assert callable(fun), "Must pass callable formula when no arity is given"
        arity = len(signature(fun).parameters)

    def wrapper(_scope: TemplateFreeScope, *indices: z3.ArithRef) -> z3.BoolRef:
        assert (
            len(indices) == arity
        ), f"Mismatch arity for formula {name}: expected {arity} got {len(indices)}"
        return fun(*indices)

    return Template(formulas={Identifier(name, arity): wrapper})


def scoped_formula(name: str, fun: _AnyScopedFormula) -> Template:
    arity = len(signature(fun).parameters) - 1

    def wrapper(scope: TemplateFreeScope, *indices: z3.ArithRef) -> z3.BoolRef:
        assert (
            len(indices) == arity
        ), f"Mismatch arity for formula {name}: expected {arity} got {len(indices)}"
        return fun(scope, *indices)

    return Template(formulas={Identifier(name, arity): wrapper})


BOUND_TRUE = bound("true", lambda _i: z3.BoolVal(True))
BOUND_LEQ_k = scoped_bound("leqk", lambda scope, i: i <= scope["k"])
BOUND_GEQ_k = scoped_bound("geqk", lambda scope, i: i >= scope["k"])

_INT_VAL_ZERO: SymTerm = lambda *indices: z3.IntVal(0)
ZERO0 = term("zero", _INT_VAL_ZERO, arity=0)
ZERO1 = term("zero", _INT_VAL_ZERO, arity=1)

INC1 = term("inc", lambda i: i + 1)
DEC1 = term("dec", lambda i: i - 1)
ID1 = term("id", lambda i: i)

SUM2 = term("sum", lambda i, j: i + j)
SUM_INC2 = term("sum-inc", lambda i, j: i + j + 1)

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

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

IS_ZERO1 = formula("is-zero", lambda i: i == 0)

IS_EVEN1 = formula("is-even", lambda i: i % 2 == 0)
IS_ODD1 = formula("is-odd", lambda i: i % 2 == 1)

IS_MOD2_REM1 = scoped_formula("is-mod2-rem", lambda scope, i: i % 2 == scope["r"])

LT2 = formula("lt", lambda i, j: i < j)
LEQ2 = formula("leq", lambda i, j: i <= j)
GT2 = formula("gt", lambda i, j: i > j)
GEQ2 = formula("geq", lambda i, j: i >= j)
EQ2 = formula("eq", lambda i, j: i == j)
EQ_INC2 = formula("eq-inc", lambda i, j: i == j + 1)
EQ_DEC2 = formula("eq-dec", lambda i, j: i == j - 1)

EQ3 = formula("eq", lambda i1, i2, i3: z3.And(i1 == i2, i1 == i3))

BETWEEN3 = formula(
    "between",
    lambda i1, i2, i3: z3.And(
        i1 <= i3,
        z3.Or(i3 < i2, i2 < i1),
    ),
)
ANTI_BETWEEN3 = formula(
    "anti-between",
    lambda i1, i2, i3: z3.And(
        i3 <= i2,
        z3.Or(i1 < i3, i2 < i1),
    ),
)

STANDARD = (
    BOUND_TRUE
    | BOUND_LEQ_k
    | BOUND_GEQ_k
    | ZERO0
    | ZERO1
    | INC1
    | DEC1
    | ID1
    | SUM2
    | SUM_INC2
    | TRUE0
    | TRUE1
    | TRUE2
    | TRUE3
    | FALSE0
    | FALSE1
    | FALSE2
    | FALSE3
    | LT2
    | LEQ2
    | GT2
    | GEQ2
    | EQ2
    | EQ_INC2
    | EQ_DEC2
    | EQ3
    | BETWEEN3
    | ANTI_BETWEEN3
)

EXTENDED = STANDARD | IS_ZERO1

SPECIAL_FORM_INT = IS_ZERO1 | TRUE1 | TRUE2 | FALSE1 | FALSE2 | EQ2


class TemplateChoice(enum.StrEnum):
    STANDARD = enum.auto()
    EXTENDED = enum.auto()

    def __call__(self) -> Template:
        match self:
            case TemplateChoice.STANDARD:
                return STANDARD
            case TemplateChoice.EXTENDED:
                return EXTENDED
