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 CallableStringTerm = Callable[[z3.SeqRef], z3.SeqRef] | Callable[
    [z3.SeqRef, z3.SeqRef], z3.ArithRef
]
type CallableStringFormula = Callable[[z3.SeqRef], z3.BoolRef] | Callable[
    [z3.SeqRef, z3.SeqRef], z3.BoolRef
] | Callable[[z3.SeqRef, z3.SeqRef, z3.SeqRef], z3.BoolRef]


class StringScope(Protocol):
    def __getitem__(self, item: str) -> z3.SeqRef: ...


def all_seq(exprs: Iterable[z3.ExprRef]) -> TypeGuard[Iterable[z3.SeqRef]]:
    return all(polymorphic_check("is_seq", expr) for expr in exprs)


def string_scope(scope: TemplateFreeScope) -> TypeGuard[StringScope]:
    return scope.theory_name == "string"


def string_bound(
    name: str, formula: Callable[[StringScope, z3.SeqRef], 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_seq(indices)
        assert string_scope(scope)
        return formula(scope, *indices)

    return TemplateFormula(name, 1, wrapper)


def string_term(name: str, term: CallableStringTerm) -> 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_seq(indices)
        return term(*indices)

    return TemplateTerm(name, arity, wrapper)


def string_formula(name: str, formula: CallableStringFormula) -> 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_seq(indices)
        return formula(*indices)

    return TemplateFormula(name, arity, wrapper)


BOUND_TRUE = string_bound("true", lambda _scope, _i: z3.BoolVal(True))

BOUND_0STAR = string_bound(
    "reg0star", lambda _scope, i: z3.InRe(i, z3.Star(z3.Re("0")))
)
BOUND_NOT_0STAR = string_bound(
    "not0star", lambda _scope, i: z3.InRe(i, z3.Complement(z3.Star(z3.Re("0"))))
)

_STRING_VAL_EMPTY = lambda: z3.StringVal("")
EMPTY0 = const_term("empty", _STRING_VAL_EMPTY, arity=0)

APPEND0_1 = string_term("append0", lambda i: i + z3.StringVal("0"))
APPEND1_1 = string_term("append1", lambda i: i + z3.StringVal("1"))
ID1 = string_term("id", lambda i: i)


IS_EMPTY1 = string_formula("is-empty", lambda i: i == z3.StringVal(""))
EQ2 = string_formula("eq", lambda i, j: i == j)
PREFIX2 = string_formula("prefix", lambda i, j: z3.PrefixOf(i, j))
FLIP_PREFIX2 = string_formula("flip-prefix", lambda i, j: z3.PrefixOf(j, i))
STRICT_PREFIX2 = string_formula(
    "strict-prefix", lambda i, j: z3.And(i != j, z3.PrefixOf(i, j))
)
IS_ZEROSTAR2 = string_formula(
    "is-zero-star", lambda i, j: z3.InRe(i, z3.Star(z3.Re("0")))
)
PREFIX_AND_DISTINCT3 = string_formula(
    "prefix-and-distinct", lambda i, j, k: z3.And(j != k, z3.PrefixOf(j, i))
)

STRING_STANDARD = Template(
    theory_name="string",
    bounds={BOUND_TRUE, BOUND_0STAR, BOUND_NOT_0STAR},
    terms={EMPTY0, ID1, APPEND0_1, APPEND1_1},
    formulas={
        TRUE0,
        TRUE1,
        TRUE2,
        TRUE3,
        FALSE0,
        FALSE1,
        FALSE2,
        FALSE3,
        EQ2,
        PREFIX2,
        FLIP_PREFIX2,
    },
)

STRING_EXTENDED = STRING_STANDARD | Template(
    theory_name="string", formulas={PREFIX_AND_DISTINCT3}
)
