from collections.abc import Callable, Iterator, MutableMapping
from dataclasses import dataclass, field, replace
from itertools import chain, product
from typing import Iterable, Self, cast

import z3

from common.relations import Relation
from common.z3_utils.rewriter import clone_vars

_COMPARISON_OPS = {z3.Z3_OP_LT, z3.Z3_OP_LE, z3.Z3_OP_GE, z3.Z3_OP_GT}


def flatten_and(expr: z3.BoolRef) -> list[z3.BoolRef]:
    if not z3.is_and(expr):
        return [expr]
    else:
        return list(
            chain(*(flatten_and(cast(z3.BoolRef, child)) for child in expr.children()))
        )


def applicable(fun: z3.FuncDeclRef, *args: z3.ExprRef) -> bool:
    if fun.arity() != len(args):
        return False

    for i, arg in enumerate(args):
        if not fun.domain(i).eq(arg.sort()):
            return False

    return True


def ground_terms(
    consts: Iterable[z3.Const], funs: Iterable[z3.FuncDeclRef], depth: int
) -> "Z3ExprSet[z3.ExprRef]":
    ground: Z3ExprSet[z3.ExprRef] = Z3ExprSet.from_iterable(consts)

    for i in range(0, depth):
        sources = ground.clone()
        new_ground: Z3ExprSet[z3.ExprRef] = Z3ExprSet()
        for fun in funs:
            for args in product(sources, repeat=fun.arity()):
                if applicable(fun, *args):
                    new_ground_term = fun(*args)
                    new_ground.add(new_ground_term)
        # new_ground = {fun(*args) for fun in funs for args in product(sources, repeat=fun.arity()) if applicable(fun, *args)}
        ground |= new_ground

    return ground


def ground_terms_by_sort(
    expr: z3.ExprRef, depth: int
) -> dict[z3.SortRef, list[z3.ExprRef]]:
    return by_sort(ground_terms_from_expr(expr, depth))


def ground_terms_from_expr(expr: z3.ExprRef, depth: int) -> Iterable[z3.ExprRef]:
    consts = constants(expr)
    funs = functions(expr)

    return ground_terms(consts, funs, depth)


def by_sort(grounds: Iterable[z3.ExprRef]) -> dict[z3.SortRef, list[z3.ExprRef]]:
    grounds_by_sort: dict[z3.SortRef, list[z3.ExprRef]] = {}
    for ground in grounds:
        sort = ground.sort()
        grounds_by_sort.setdefault(sort, [])
        grounds_by_sort[sort].append(ground)

    return grounds_by_sort


def ground_terms_by_sort_non_repeating(
    expr: z3.ExprRef,
) -> dict[z3.SortRef, list[z3.ExprRef]]:
    consts = constants(expr)
    funs = functions(expr)

    return by_sort(ground_terms_non_repeating(consts, funs))


@dataclass
class FunsPerTerm:
    data: dict[z3.ExprRef, frozenset[z3.FuncDeclRef]] = field(default_factory=dict)

    def __getitem__(self, item: z3.ExprRef) -> frozenset[z3.FuncDeclRef]:
        if item not in self.data:
            self.data[item] = functions(item)

        return self.data[item]


def ground_terms_non_repeating(
    consts: frozenset[z3.Const], funs: frozenset[z3.FuncDeclRef]
) -> Iterable[z3.ExprRef]:
    funs_per_term = FunsPerTerm()

    old_terms: set[z3.ExprRef] = set()
    new_terms: set[z3.ExprRef] = {const for const in consts}

    while not (old_terms >= new_terms):
        old_terms.update(new_terms)
        new_terms = {
            fun(*args)
            for fun in funs
            for args in product(old_terms, repeat=fun.arity())
            if applicable(fun, *args)
            and not any(fun in funs_per_term[arg] for arg in args)
        }

    return old_terms


def is_bool_value(expression: z3.ExprRef) -> bool:
    return z3.is_true(expression) or z3.is_false(expression)


def constants(
    expression: z3.ExprRef, sort: z3.SortRef | None = None
) -> frozenset[z3.Const]:
    if (
        z3.is_const(expression)
        and not z3.is_var(expression)
        and not z3.is_int_value(expression)
        and not is_bool_value(expression)
        and not z3.is_or(expression)
        and not z3.is_and(expression)
        and (sort is None or expression.sort() == sort)
    ):
        return frozenset({expression})

    return frozenset(
        {const for child in expression.children() for const in constants(child, sort)}
    )


def uninterpreted_constants(
    expression: z3.ExprRef, sort: z3.SortRef | None = None
) -> set[z3.Const]:
    if (
        z3.is_const(expression)
        and expression.decl().kind() == z3.Z3_OP_UNINTERPRETED
        and (
            (sort is None and expression.sort().kind() == z3.Z3_UNINTERPRETED_SORT)
            or expression.sort() == sort
        )
    ):
        return {expression}

    return {
        const
        for child in expression.children()
        for const in uninterpreted_constants(child, sort)
    }


def sorts_of_decl(decl: z3.FuncDeclRef) -> set[z3.SortRef]:
    return {decl.range()} | {decl.domain(i) for i in range(decl.arity())}


def uninterpreted_sorts_of_decl(decl: z3.FuncDeclRef) -> set[z3.SortRef]:
    return {
        sort for sort in sorts_of_decl(decl) if sort.kind() == z3.Z3_UNINTERPRETED_SORT
    }


def uninterpreted_sorts(expression: z3.ExprRef) -> set[z3.SortRef]:
    sorts = set()
    if z3.is_app(expression):
        sorts = uninterpreted_sorts_of_decl(expression.decl())
    return sorts | {
        sort for child in expression.children() for sort in uninterpreted_sorts(child)
    }


def functions(expression: z3.ExprRef) -> frozenset[z3.FuncDeclRef]:
    sub_functions = {fun for child in expression.children() for fun in functions(child)}
    if (
        z3.is_app(expression)
        and not z3.is_bool(expression)
        and not z3.is_const(expression)
    ):
        sub_functions |= {expression.decl()}

    return frozenset(sub_functions)


def uninterpreted_functions(expression: z3.ExprRef) -> set[z3.FuncDeclRef]:
    sub_functions = {
        fun for child in expression.children() for fun in uninterpreted_functions(child)
    }
    if (
        z3.is_app(expression)
        and not z3.is_bool(expression)
        and not z3.is_const(expression)
        and expression.decl().kind() == z3.Z3_OP_UNINTERPRETED
    ):
        sub_functions |= {expression.decl()}

    return sub_functions


def uninterpreted_relations(expression: z3.ExprRef) -> set[Relation]:
    sub_relations = {
        fun for child in expression.children() for fun in uninterpreted_relations(child)
    }

    if (
        z3.is_app(expression)
        and z3.is_bool(expression)
        and not z3.is_const(expression)
        and expression.decl().kind() == z3.Z3_OP_UNINTERPRETED
    ):
        sub_relations |= {cast(Relation, expression.decl())}

    return sub_relations


def signature(f: z3.FuncDeclRef) -> Iterable[z3.SortRef]:
    return tuple(domain(f)) + (f.range(),)


def domain(f: z3.FuncDeclRef) -> Iterable[z3.SortRef]:
    return tuple(f.domain(i) for i in range(f.arity()))


def terms(expression: z3.ExprRef) -> frozenset[z3.ExprRef]:
    sub_terms = frozenset(
        {term for child in expression.children() for term in terms(child)}
    )

    if z3.is_app(expression) and not z3.is_bool(expression):
        sub_terms |= {expression}

    return sub_terms


def decls(expression: z3.ExprRef) -> frozenset[z3.FuncDeclRef]:
    return decl(expression).union(*[decls(c) for c in expression.children()])


def decl(expression: z3.ExprRef) -> frozenset[z3.FuncDeclRef]:
    if z3.is_app(expression):
        return frozenset({expression.decl()})
    else:
        return frozenset()


def unpack_quantifier(
    quantifier: z3.QuantifierRef,
) -> tuple[list[z3.Const], z3.BoolRef]:
    bounding_vars = clone_vars(quantifier)

    body = z3.substitute_vars(
        quantifier.body(), *reversed(bounding_vars)
    )  # Z3 uses vars in reverse order

    return bounding_vars, body


def transform_quantifier(
    quantifier: z3.QuantifierRef, transformer: Callable[[z3.BoolRef], z3.BoolRef]
) -> z3.QuantifierRef:
    variables, body = unpack_quantifier(quantifier)

    new_body = transformer(body)

    qid = quantifier.qid()
    if qid == "<null>":
        qid = ""

    if quantifier.is_exists():
        return z3.Exists(variables, new_body, qid=qid)
    else:
        return z3.ForAll(variables, new_body, qid=qid)


@dataclass
class Z3ExprWrapper:
    expr: z3.ExprRef

    def __hash__(self) -> int:
        return hash(self.expr)

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Z3ExprWrapper):
            return False

        return self.expr.eq(other.expr)


@dataclass
class Z3ExprSet[Ref: z3.ExprRef]:
    elements: set[Z3ExprWrapper] = field(default_factory=set)

    def clone(self) -> Self:
        return replace(self, elements=set(self.elements))

    @classmethod
    def from_iterable(cls, iterable: Iterable[Ref]) -> Self:
        return cls({Z3ExprWrapper(expr) for expr in iterable})

    @classmethod
    def from_args(cls, *args: Ref) -> Self:
        return cls.from_iterable(args)

    def add(self, value: Ref) -> None:
        self.elements.add(Z3ExprWrapper(value))

    def discard(self, value: z3.ExprRef) -> None:
        self.elements.discard(Z3ExprWrapper(value))

    def __contains__(self, x: object) -> bool:
        if not isinstance(x, z3.ExprRef):
            return False

        return Z3ExprWrapper(x) in self.elements

    def __len__(self) -> int:
        return len(self.elements)

    def __iter__(self) -> Iterator[z3.ExprRef]:
        return iter([element.expr for element in self.elements])

    def __or__(self, other: Self) -> Self:
        return replace(self, elements=self.elements | other.elements)


@dataclass
class Z3ExprDict[Ref: z3.ExprRef, T](MutableMapping[Ref, T]):
    mapping: dict[Z3ExprWrapper, T] = field(default_factory=dict)

    @classmethod
    def from_pairs(cls, *pairs: tuple[Ref, T]) -> Self:
        return cls({Z3ExprWrapper(key): value for key, value in pairs})

    def __setitem__(self, key: Ref, value: T, /) -> None:
        self.mapping[Z3ExprWrapper(key)] = value

    def __delitem__(self, key: Ref, /) -> None:
        del self.mapping[Z3ExprWrapper(key)]

    def __getitem__(self, key: Ref, /) -> T:
        return self.mapping[Z3ExprWrapper(key)]

    def __len__(self) -> int:
        return len(self.mapping)

    def __iter__(self) -> Iterator[Ref]:
        return iter([cast(Ref, element.expr) for element in self.mapping])


def ground_terms_in_formula(formula: z3.ExprRef) -> Z3ExprSet[z3.ExprRef]:
    the_ground_terms: Z3ExprSet[z3.ExprRef] = Z3ExprSet()

    def rec_ground_terms_in_expr(expr: z3.ExprRef) -> bool:
        if z3.is_const(expr):
            the_ground_terms.add(expr)
            return True

        all_children_ground = True
        for child in expr.children():
            if not rec_ground_terms_in_expr(child):
                all_children_ground = False

        if (
            z3.is_app(expr)
            and not z3.is_bool(expr)
            and not z3.is_const(expr)
            and expr.decl().kind() == z3.Z3_OP_UNINTERPRETED
            and all_children_ground
        ):
            the_ground_terms.add(expr)
            return True

        return False

    rec_ground_terms_in_expr(formula)

    return the_ground_terms


def ground_terms_in_formula_by_sort(
    formula: z3.ExprRef,
) -> dict[z3.SortRef, list[z3.ExprRef]]:
    return by_sort(iter(ground_terms_in_formula(formula)))


def is_comparison(formula: z3.BoolRef) -> bool:
    return formula.decl().kind() in _COMPARISON_OPS


def is_empty_and(formula: z3.ExprRef) -> bool:
    return z3.is_const(formula) and formula.decl().name() == "and"


def is_empty_or(formula: z3.ExprRef) -> bool:
    return z3.is_const(formula) and formula.decl().name() == "or"


def is_universal(formula: z3.ExprRef, negated: bool = False) -> bool:
    if z3.is_quantifier(formula):
        if negated == formula.is_forall():
            return False
    elif z3.is_not(formula):
        return is_universal(formula.children()[0], not negated)
    elif z3.is_implies(formula):
        antecedent, consequent = formula.children()
        return is_universal(antecedent, not negated) and is_universal(
            consequent, negated
        )

    return all(map(lambda child: is_universal(child, negated), formula.children()))


def exactly_one(*exprs: z3.BoolRef) -> z3.BoolRef:
    if len(exprs) == 0:
        raise RuntimeError("Zero expressions given")
    if len(exprs) == 1:
        return exprs[0]
    elif len(exprs) == 2:
        a, b = exprs
        return z3.Xor(a, b)

    options = []
    for i, expr in enumerate(exprs):
        option = []
        for j, expr2 in enumerate(exprs):
            if i == j:
                option.append(expr2)
            else:
                option.append(z3.Not(expr2))
        options.append(z3.And(*option))
    return z3.Or(*options)

    # return (
    #     reduce(
    #         operator.add,
    #         (z3.If(expr, z3.IntVal(1), z3.IntVal(0)) for expr in exprs),
    #     )
    #     == 1
    # )
