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

import cvc5.pythonic as cvc5
import z3
from pymona.strings.z3_api import PyMONAZ3Solver

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 (
    constants,
    is_comparison,
    uninterpreted_functions,
    unpack_quantifier,
)
from fest import Backend
from fest.cvc5_utils import cvc5_solver_for
from fest.theories import BaseTheory, TheoryModel, TheoryName, get_theory
from fest.z3_utils import well_definedness_formula

from . import Assignment, Node, SymElement, SymStruct


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

    @cached_property
    def theory_name(self) -> TheoryName:
        return self.struct.theory_name

    @cached_property
    def theory(self) -> BaseTheory:
        return self.struct.theory

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

    @cached_property
    def well_definedness_formula(self) -> z3.BoolRef:
        return self._transform_formula(
            well_definedness_formula(
                constants(self.formula), uninterpreted_functions(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(),
                self.theory_name,
                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 = [sorted(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,
                    )
                )

        base_theory_variables = [self._base_const(var) for var in bounding_vars]

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

    def _transform_children(
        self, expr: z3.ExprRef, assignment: Assignment
    ) -> tuple[Iterable[Node], Iterable[z3.ExprRef]]:
        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(), self.theory_name))
            else:
                index = self._base_const(var)
                guards.append(bounds(index))
                elements.append(SymElement(node, var.sort(), self.theory_name, index))

        return guards, elements

    def _base_const(self, const: z3.Const) -> z3.Const:
        name = const.decl().name()
        return self.theory.to_base_const(name)


def check_sym_struct(
    formula: z3.BoolRef, struct: SymStruct, assignment: Assignment | None = None
) -> z3.CheckSatResult:
    if assignment is None:
        assignment = Assignment()

    result = resolve_assignment(formula, struct, assignment)
    if result == "unsat":
        return z3.unsat
    elif result == "unknown":
        return z3.unknown
    else:
        return z3.sat


type BaseQueryResult = Literal["unknown", "unsat"] | TheoryModel


@config(backend=config.fest.backend[str])
def resolve_assignment(
    formula: z3.BoolRef,
    struct: SymStruct,
    assignment: Assignment,
    *,
    backend: Backend = "auto",
) -> BaseQueryResult:
    if assignment is None:
        assignment = Assignment()
    checker = ModelChecker(formula, struct, assignment)

    result = _check_base_query(
        checker.well_definedness_formula, struct.theory_name, backend
    )
    if result == "unsat" or result == "unknown":
        raise ValueError("Structure is not well-defined")

    return _check_base_query(checker.result, struct.theory_name, backend)


def _check_base_query(
    formula: z3.BoolRef, theory_name: TheoryName, backend: Backend
) -> BaseQueryResult:
    theory = get_theory(theory_name)
    if backend == "cvc5" or (backend == "auto" and theory_name == "lia"):
        cvc5_solver = cvc5_solver_for(theory.logic())
        result = cvc5_solver.check(convert_z3_to_cvc5(formula))
        if result == cvc5.sat:
            return theory.get_model(cvc5_solver.model())
        elif result == cvc5.unsat:
            return "unsat"
        else:
            return "unknown"
    elif backend == "mona" or (backend == "auto" and theory_name == "string"):
        mona_solver = PyMONAZ3Solver()
        mona_result = mona_solver.check(formula)
        if mona_result == z3.sat:
            return theory.get_model(mona_solver.model())
        elif mona_result == z3.unsat:
            return "unsat"
        else:
            return "unknown"
    else:
        z3_solver = z3.SolverFor(theory.logic())
        z3_result = z3_solver.check(formula)
        if z3_result == z3.sat:
            return theory.get_model(z3_solver.model())
        elif z3_result == z3.unsat:
            return "unsat"
        else:
            return "unknown"
