from collections.abc import Iterable
from dataclasses import dataclass
from functools import cached_property
from itertools import product
from typing import cast

import cvc5.pythonic as cvc5
import z3

from common import config
from common.converters import convert_z3_to_cvc5
from common.relations import is_relation
from common.z3_utils.expr_utils import is_comparison, unpack_quantifier
from fest.cvc5_utils import cvc5_lia_solver
from fest.z3_utils import to_lia_variable

from . import Assignment, Node, SymElement, SymStruct


@dataclass(frozen=True)
class Transformer:
    formula: z3.BoolRef
    struct: SymStruct
    assignment: Assignment

    @cached_property
    def result(self) -> z3.BoolRef:
        return self._transform_formula(self.formula, self.assignment)

    def _transform_formula(
        self, formula: z3.BoolRef, assignment: Assignment
    ) -> z3.BoolRef:
        if z3.is_true(formula) or z3.is_false(formula):
            return formula
        elif z3.is_quantifier(formula):
            return self._transform_quantifier(formula, assignment)
        elif formula.decl().kind() == z3.Z3_OP_UNINTERPRETED:
            return self._transform_relation(formula, assignment)
        elif z3.is_eq(formula):
            return self._transform_equality(formula, assignment)
        elif z3.is_distinct(formula):
            return z3.Not(self._transform_equality(formula, assignment))
        elif is_comparison(formula):
            assert config.fest.experimental.int_support(
                bool
            ), "`fest.experimental.int_support` is not enabled"
            return formula
        else:
            form = str(formula.decl())
            children = [
                self._transform_formula(cast(z3.BoolRef, child), assignment)
                for child in formula.children()
            ]

            return cast(z3.BoolRef, getattr(z3, form)(*children))

    def _transform_relation(
        self, formula: z3.BoolRef, assignment: Assignment
    ) -> z3.BoolRef:
        relation = formula.decl()

        assert is_relation(relation), f"Function used as relation {relation}"

        nodes, indices = self._transform_children(formula, assignment)

        return self.struct[("rel", relation, *nodes)](*indices)

    def _transform_equality(
        self, formula: z3.BoolRef, assignment: Assignment
    ) -> z3.BoolRef:
        if z3.is_bool(formula.children()[0]):  # iff construction
            left_formula, right_formula = [
                self._transform_formula(cast(z3.BoolRef, child), assignment)
                for child in formula.children()
            ]
            return left_formula == right_formula

        (left, right), (i, j) = self._transform_children(formula, assignment)

        if left != right:
            return z3.BoolVal(False)
        else:
            return i == j

    def _transform_term(self, term: z3.ExprRef, assignment: Assignment) -> SymElement:
        if z3.is_int(term):
            assert config.fest.experimental.int_support(
                bool
            ), "`fest.experimental.int_support` is not enabled"
            return SymElement.int(term)

        if z3.is_const(term):
            name = str(term)
            if name in assignment:
                return assignment[str(term)]
            elif term in self.struct.constants:
                return self.struct.constants[term]
            else:
                assert False, f"Missing interpretation for `{name}`"

        fun = term.decl()
        if fun.kind() == z3.Z3_OP_UNINTERPRETED:
            nodes, indices = self._transform_children(term, assignment)

            node, sym_term = self.struct[("fun", fun, *nodes)]

            index = sym_term(*indices)

            return SymElement(node, fun.range(), index)

        assert False, f"Unsupported term {term}"

    def _transform_quantifier(
        self, quantifier: z3.QuantifierRef, assignment: Assignment
    ) -> z3.BoolRef:
        bounding_vars, body = unpack_quantifier(quantifier)

        options = [self.struct.domain[var.sort()] for var in bounding_vars]

        bodies = []

        for option in product(*options):
            bounds, elements = self._get_bounded_elements(bounding_vars, option)

            transformed_body = self._transform_formula(
                body, assignment.override(bounding_vars, elements)
            )

            if quantifier.is_exists():
                bodies.append(
                    z3.And(
                        *bounds,
                        transformed_body,
                    )
                )
            else:
                bodies.append(
                    z3.Implies(
                        z3.And(*bounds),
                        transformed_body,
                    )
                )

        lia_variables = [to_lia_variable(var) for var in bounding_vars]

        if quantifier.is_exists():
            return z3.Exists(lia_variables, z3.Or(*bodies))
        else:
            return z3.ForAll(lia_variables, z3.And(*bodies))

    def _transform_children(
        self, expr: z3.ExprRef, assignment: Assignment
    ) -> tuple[Iterable[Node], Iterable[z3.ArithRef]]:
        children = [
            self._transform_term(child, assignment) for child in expr.children()
        ]

        nodes = (element.node for element in children)
        indices = (element.to_index() for element in children)

        return nodes, indices

    def _get_bounded_elements(
        self, variables: list[z3.Const], nodes: tuple[Node, ...]
    ) -> tuple[list[z3.BoolRef], list[SymElement]]:
        guards = []
        elements = []
        for var, node in zip(variables, nodes):
            bounds = self.struct[node]
            if bounds is None:
                elements.append(SymElement(node, var.sort()))
            else:
                lia_variable = to_lia_variable(var)
                guards.append(bounds(lia_variable))
                elements.append(SymElement(node, var.sort(), lia_variable))

        return guards, elements


@config(use_cvc5=config.fest.cvc5)
def check_sym_struct(
    formula: z3.BoolRef,
    struct: SymStruct,
    assignment: Assignment | None = None,
    *,
    use_cvc5: bool = False,
) -> z3.CheckSatResult:
    if assignment is None:
        assignment = Assignment()
    lia_query = Transformer(formula, struct, assignment).result
    if use_cvc5:
        result = cvc5_lia_solver().check(convert_z3_to_cvc5(lia_query))
        if result == cvc5.sat:
            return z3.sat
        elif result == cvc5.unsat:
            return z3.unsat
        else:
            return z3.unknown
    else:
        return z3.Solver().check(lia_query)
