import enum
from collections.abc import Callable, Iterable
from inspect import signature
from typing import Protocol, TypeGuard

import z3

from common.z3_utils.patcher import polymorphic_check
from fest.templates import Template, TemplateFormula, TemplateTerm
from fest.templates.template_free_variables import TemplateFreeScope

from . import FALSE0, FALSE1, FALSE2, FALSE3, TRUE0, TRUE1, TRUE2, TRUE3, const_term

type CallableArithTerm = Callable[[z3.ArithRef], z3.ArithRef] | Callable[
    [z3.ArithRef, z3.ArithRef], z3.ArithRef
]
type CallableArithFormula = Callable[[z3.ArithRef], z3.BoolRef] | Callable[
    [z3.ArithRef, z3.ArithRef], z3.BoolRef
] | Callable[[z3.ArithRef, z3.ArithRef, z3.ArithRef], z3.BoolRef] | Callable[
    [z3.ArithRef, z3.ArithRef, z3.ArithRef, z3.ArithRef], z3.BoolRef
]


class ArithScope(Protocol):
    def __getitem__(self, item: str) -> z3.ArithRef: ...


def all_arith(exprs: Iterable[z3.ExprRef]) -> TypeGuard[Iterable[z3.ArithRef]]:
    return all(polymorphic_check("is_arith", expr) for expr in exprs)


def arith_scope(scope: TemplateFreeScope) -> TypeGuard[ArithScope]:
    return scope.theory_name == "lia"


def arith_scoped_formula(
    name: str, formula: Callable[[ArithScope, z3.ArithRef], z3.BoolRef]
) -> TemplateFormula:
    def wrapper(scope: TemplateFreeScope, *indices: z3.ExprRef) -> z3.BoolRef:
        assert (
            len(indices) == 1
        ), f"Mismatch arity for bound formula {name}: expected 1 got {len(indices)}"
        assert all_arith(indices)
        assert arith_scope(scope)
        return formula(scope, *indices)

    return TemplateFormula(name, 1, wrapper)


def arith_term(name: str, term: CallableArithTerm) -> TemplateTerm:
    arity = len(signature(term).parameters)

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

    return TemplateTerm(name, arity, wrapper)


def arith_formula(name: str, formula: CallableArithFormula) -> TemplateFormula:
    arity = len(signature(formula).parameters)

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

    return TemplateFormula(name, arity, wrapper)


LEQ_k = arith_scoped_formula("leqk", lambda scope, i: i <= scope["k"])
GEQ_k = arith_scoped_formula("geqk", lambda scope, i: i >= scope["k"])

_INT_VAL_ZERO = lambda: z3.IntVal(0)
ZERO0 = const_term("zero", _INT_VAL_ZERO, arity=0)

INC1 = arith_term("inc", lambda i: i + 1)
ZERO1 = const_term("zero", _INT_VAL_ZERO, arity=1)
DEC1 = arith_term("dec", lambda i: i - 1)
ID1 = arith_term("id", lambda i: i)

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


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

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

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

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

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


def is_diff_odd(x: z3.ArithRef, y: z3.ArithRef) -> z3.BoolRef:
    return z3.And(x <= y, (y - x) % 2 == 1)


def is_diff_even(x: z3.ArithRef, y: z3.ArithRef) -> z3.BoolRef:
    return z3.And(x <= y, (y - x) % 2 == 0)


def is_diff(x: z3.ArithRef, y: z3.ArithRef) -> z3.BoolRef:
    return x <= y


def between(x: z3.ArithRef, y: z3.ArithRef, e: z3.ArithRef) -> z3.BoolRef:
    return z3.And(x <= e, e < y)


IS_DIFF_ODD2 = arith_formula("is-diff-odd", is_diff_odd)
IS_DIFF_EVEN2 = arith_formula("is-diff-even", is_diff_even)
BETWEEN_DIFF3 = arith_formula(
    "between-diff",
    lambda i, j, k: z3.Or(
        z3.And(between(i, j, k), is_diff(i, j)), z3.Not(is_diff(i, j))
    ),
)
BETWEEN_DIFF_ODD3 = arith_formula(
    "between-diff-odd",
    lambda i, j, k: z3.Or(
        z3.And(between(i, j, k), is_diff_even(i, j)), z3.Not(is_diff_even(i, j))
    ),
)
BETWEEN_DIFF_EVEN3 = arith_formula(
    "between-diff-even",
    lambda i, j, k: z3.Or(
        z3.And(between(i, j, k), is_diff_odd(i, j)), z3.Not(is_diff_odd(i, j))
    ),
)
LEQ_IGNORE_MIDDLE3 = arith_formula("leq-ignore-middle", lambda i, j, k: i <= k)
LEQ_IGNORE_LAST3 = arith_formula("leq-ignore-last", lambda i, j, k: i <= j)
BETWEEN_DIFF_IGNORE_THIRD4 = arith_formula(
    "between-diff-ignore-third",
    lambda i, j, k, l: z3.Or(z3.And(i <= l, l < j, i <= j), z3.Not(i <= j)),
)
SIMPLE_BETWEEN3 = arith_formula("simple-between", between)
SIMPLE_BETWEEN_IGNORE_THIRD4 = arith_formula(
    "simple-between-ignore-third", lambda i, j, k, w: z3.And(i <= w, w < j)
)

EQ_DIFF3 = arith_formula("eq-diff", lambda i, j, k: z3.And(i <= j, k == j - i))

STANDARD = Template(
    bounds={TRUE1, LEQ_k, GEQ_k},
    terms={ZERO0, ZERO1, INC1, DEC1, ID1, SUM2, SUM_INC2},
    formulas={
        TRUE0,
        TRUE1,
        TRUE2,
        TRUE3,
        FALSE0,
        FALSE1,
        FALSE2,
        FALSE3,
        LT2,
        LEQ2,
        GT2,
        GEQ2,
        EQ2,
        EQ_INC2,
        EQ_DEC2,
        EQ3,
        BETWEEN3,
        ANTI_BETWEEN3,
    },
)
EXTENDED = STANDARD | Template(formulas={IS_ZERO1})
SLS = Template(
    bounds={TRUE1},
    terms={ZERO0, INC1},
    formulas={
        FALSE0,
        FALSE1,
        FALSE2,
        FALSE3,
        TRUE0,
        TRUE1,
        TRUE2,
        TRUE3,
        LEQ2,
        IS_DIFF_EVEN2,
        IS_DIFF_ODD2,
        BETWEEN_DIFF3,
        BETWEEN_DIFF_EVEN3,
        BETWEEN_DIFF_ODD3,
        BETWEEN3,
    },
)

_SPECIAL_FORM_INT = Template(
    formulas={
        GEQ_k,
        LEQ_k,
        IS_ZERO1,
        TRUE1,
        TRUE2,
        FALSE1,
        FALSE2,
        EQ2,
        LEQ_IGNORE_MIDDLE3,
    }
)

SPECIAL_FORM_INT = (
    Template(
        formulas={
            IS_ZERO1,
            TRUE1,
            TRUE2,
            TRUE3,
            FALSE1,
            FALSE2,
            FALSE3,
            EQ_DIFF3,
            LEQ_IGNORE_MIDDLE3,
            SIMPLE_BETWEEN_IGNORE_THIRD4,
        }
    )
    | _SPECIAL_FORM_INT
)


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

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