import io
import operator
from dataclasses import dataclass
from functools import cached_property, reduce
from itertools import product
from typing import 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 Relation
from common.simplifier import Simplifier
from common.z3_utils.expr_utils import (
    constants,
    exactly_one,
    is_comparison,
    uninterpreted_constants,
    uninterpreted_functions,
    unpack_quantifier,
)
from common.z3_utils.patcher import is_z3_as_cvc5, z3_as_cvc5, z3_as_z3
from fest import Backend
from fest.cvc5_utils import cvc5_solver_for
from fest.symstructs import Assignment, SymElement, SymStruct
from fest.theories import TheoryCacheDelegate, TheoryModel, TheoryNamed
from fest.z3_utils import find_order_relation, well_definedness_formula

from . import Template
from .guards import Clause, Generator
from .models import Model
from .stdlib.arith import SPECIAL_FORM_INT
from .template_free_variables import TemplateFreeVariablesGenerator
from .utils import lexicographical


def _convert_bool(formula: z3.BoolRef) -> z3.BoolRef:
    if is_z3_as_cvc5():
        return cast(z3.BoolRef, convert_z3_to_cvc5(formula))
    else:
        return formula


def _convert_arith(term: z3.ArithRef) -> z3.ArithRef:
    if is_z3_as_cvc5():
        return cast(z3.ArithRef, convert_z3_to_cvc5(term))
    else:
        return term


@dataclass(frozen=True)
class ModelFinder(TheoryCacheDelegate):
    formula: z3.BoolRef
    template: Template

    @cached_property
    def delegate(self) -> TheoryNamed:
        return self.template

    @cached_property
    def simplifier(self) -> Simplifier:
        return Simplifier(config.fest.simplify(bool))

    @cached_property
    def true(self) -> z3.BoolRef:
        return z3.BoolVal(True)

    @cached_property
    def false(self) -> z3.BoolRef:
        return z3.BoolVal(False)

    @cached_property
    def domain(self) -> dict[z3.SortRef, set[SymElement]]:
        index: z3.ExprRef = self.theory.to_base_const("i")
        domain = {
            sort: {
                SymElement(f"{sort}_{i}", sort, self.theory_name)
                for i in range(domain.regular)
            }
            | {
                SymElement(
                    f"{sort}_{i + domain.regular}",
                    sort,
                    self.theory_name,
                    index,
                )
                for i in range(domain.summary)
            }
            for sort, domain in self.template.domain.items()
        }

        if config.fest.experimental.int_support(bool):
            domain[z3.IntSort()] = {SymElement.int()}

        return domain

    @cached_property
    def constants(self) -> dict[z3.SortRef, list[z3.Const]]:
        return {
            sort: list(constants(self.formula, sort))
            for sort in self.domain
            if not sort.eq(z3.IntSort())
        }

    @cached_property
    def functions(self) -> set[z3.FuncDeclRef]:
        return uninterpreted_functions(self.formula)

    @cached_property
    def order(self) -> Relation | None:
        return find_order_relation(self.formula)

    @cached_property
    def first_constant_of_sort(self) -> dict[z3.SortRef, z3.Const]:
        return {
            sort: sort_constants[0]
            for sort, sort_constants in self.constants.items()
            if len(sort_constants) > 0
        }

    @cached_property
    def default_assignment(self) -> Assignment:
        if (
            config.fest.experimental.allow_summary_constants
            and not config.fest.experimental.assume_regular_first_constant
        ):
            return Assignment()

        return Assignment(
            {
                str(constant): SymElement(f"{sort}_0", sort, self.theory_name)
                for sort, constant in self.first_constant_of_sort.items()
            }
        )

    @cached_property
    def guards(self) -> Generator:
        return Generator()

    @cached_property
    def variables(self) -> TemplateFreeVariablesGenerator:
        return TemplateFreeVariablesGenerator(self.theory)

    @cached_property
    def result(self) -> z3.BoolRef:
        base_theory_query = self._transform_formula(
            self.formula, self.default_assignment
        )

        return self.simplifier.simplified_and(
            base_theory_query, self._generate_aux(self.default_assignment)
        )

    def _transform_formula(
        self, formula: z3.BoolRef, assignment: Assignment
    ) -> z3.BoolRef:
        if z3.is_true(formula):
            return self.true
        elif z3.is_false(formula):
            return self.false
        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 self.simplifier.simplified_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 _convert_bool(formula)
        else:
            form = str(formula.decl())
            children = [
                self._transform_formula(cast(z3.BoolRef, child), assignment)
                for child in formula.children()
            ]

            if z3.is_and(formula):
                return self.simplifier.simplified_and(*children)
            elif z3.is_or(formula):
                return self.simplifier.simplified_or(*children)
            elif z3.is_not(formula):
                return self.simplifier.simplified_not(children[0])
            elif z3.is_implies(formula):
                return self.simplifier.simplified_implies(*children)

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

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

        if self.order is not None and self.order.eq(relation):
            return self._transform_order(formula, assignment)

        children = [
            self._transform_term(child, assignment).items()
            for child in formula.children()
        ]

        options = {}

        for option in product(*children):
            clause, elements = _transpose_option(*option)

            if clause is None:
                continue

            nodes = ", ".join(element.node for element in elements)
            indices = [
                element.to_index() for element in elements if element.is_summary()
            ]

            if config.fest.experimental.regular_ints:
                indices = [
                    element.to_index()
                    for element in elements
                    if element.is_summary() and not element.sort.eq(z3.IntSort())
                ]

            for guard, result in (
                self._get_template(relation)
                .apply_formulas(
                    *indices, scope=self.variables.scope(f"{relation}({nodes})")
                )
                .items()
            ):
                if (
                    clause_with_relation_guard := clause.add(
                        self.guards[f"{relation}({nodes})", guard]
                    )
                ) is not None:
                    options[clause_with_relation_guard] = result

        return self.simplifier.simplified_and(
            *(
                self.simplifier.simplified_implies(key.to_bool(self.simplifier), value)
                for key, value in options.items()
            )
        )

    def _transform_order(
        self, formula: z3.BoolRef, assignment: Assignment
    ) -> z3.BoolRef:
        children = [
            self._transform_term(child, assignment).items()
            for child in formula.children()
        ]

        options = {}

        for option in product(*children):
            clause, elements = _transpose_option(*option)

            if clause is None:
                continue

            left, right = elements

            is_lexicographical, first, second = lexicographical(left.sort, left, right)
            nodes = f"{first.node}, {second.node}"

            for guard, result in self._order_formulas(first, second):
                if (
                    clause_with_relation_guard := clause.add(
                        self.guards[f"{self.order}({nodes})", guard]
                    )
                ) is not None:
                    if not is_lexicographical:
                        result = self.simplifier.simplified_not(result)

                    options[clause_with_relation_guard] = result

        return self.simplifier.simplified_and(
            *(
                self.simplifier.simplified_implies(key.to_bool(self.simplifier), value)
                for key, value in options.items()
            )
        )

    def _order_formulas(
        self, first: SymElement, second: SymElement
    ) -> list[tuple[str, z3.BoolRef]]:
        assert self.theory.name() == "lia", "Order only supported in LIA for now"
        i1, i2 = first.to_index(), second.to_index()
        any_regular = first.is_regular() or second.is_regular()
        same_node = first.node == second.node

        if any_regular and same_node:
            return [("bottom", self.false)]
        elif any_regular:  # not same_node
            return [("top", self.true), ("bottom", self.false)]
        elif same_node:  # not any_regular
            return [("lt", cast(z3.ArithRef, i1) < cast(z3.ArithRef, i2))]
        else:  # not any_regular and not same_node
            return [
                ("top", self.true),
                ("leq", cast(z3.ArithRef, i1) <= cast(z3.ArithRef, i2)),
                ("bottom", self.false),
            ]

    def _transform_equality(
        self, formula: z3.BoolRef, assignment: Assignment
    ) -> z3.BoolRef:
        assert len(formula.children()) == 2, "Non-binary equality unsupported"

        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 self.simplifier.simplified_iff(left_formula, right_formula)

        # term equality
        children = [
            self._transform_term(child, assignment).items()
            for child in formula.children()
        ]

        options = {}

        for option in product(*children):
            clause, (left, right) = _transpose_option(*option)

            if clause is None:
                continue

            if left.node != right.node:
                options[clause] = self.false
            elif left.is_regular():
                options[clause] = self.true
            else:
                formula1 = left.to_index()
                formula2 = right.to_index()
                options[clause] = formula1 == formula2

        return self.simplifier.simplified_and(
            *(
                self.simplifier.simplified_implies(key.to_bool(self.simplifier), value)
                for key, value in options.items()
            )
        )

    def _transform_term(
        self, term: z3.ExprRef, assignment: Assignment
    ) -> dict[Clause, SymElement]:
        if z3.is_int(term):
            assert config.fest.experimental.int_support(
                bool
            ), "`fest.experimental.int_support` is not enabled"
            return {Clause(): SymElement.int(_convert_arith(term))}
        elif z3.is_const(term):
            return self._transform_const(term, assignment)
        elif term.decl().kind() == z3.Z3_OP_UNINTERPRETED:
            return self._transform_fun(term, assignment)

        assert False, f"Unsupported term {term}"

    def _transform_fun(
        self, term: z3.ExprRef, assignment: Assignment
    ) -> dict[Clause, SymElement]:
        fun = term.decl()
        children = [
            self._transform_term(child, assignment).items() for child in term.children()
        ]
        options = {}
        for option in product(*children):
            clause, elements = _transpose_option(*option)

            if clause is None:
                continue

            nodes = ", ".join(element.node for element in elements)

            for target in self.domain[fun.range()]:
                if (
                    clause_with_target_guard := clause.add(
                        self.guards[f"{fun}({nodes})->", f"{target.node}"]
                    )
                ) is not None:
                    if target.is_regular():
                        options[clause_with_target_guard] = target
                    else:
                        indices = [
                            element.to_index()
                            for element in elements
                            if element.is_summary()
                        ]

                        for guard, result in self.template.apply_terms(
                            *indices,
                            scope=self.variables.scope(f"{fun}({nodes})"),
                        ).items():
                            if (
                                clause_with_term_guard := clause_with_target_guard.add(
                                    self.guards[f"{fun}({nodes})", guard]
                                )
                            ) is not None:
                                options[clause_with_term_guard] = target.with_index(
                                    result
                                )

        assert options, f"No available function options for `{term}`"

        return options

    def _transform_const(
        self, term: z3.Const, assignment: Assignment
    ) -> dict[Clause, SymElement]:
        const = term.decl().name()
        value = assignment.get(const, None)
        if value is not None:
            return {Clause(): value}
        else:
            options = {}
            for element in self.domain[term.sort()]:
                if (
                    element.is_summary()
                    and not config.fest.experimental.allow_summary_constants
                ):
                    continue

                clause = Clause.create(self.guards[const, element.node])
                if element.is_regular():
                    options[clause] = element
                else:
                    index = self.variables.scope(const)["i"]
                    options[clause] = element.with_index(index)
            return options

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

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

        bodies = []

        for option in product(*options):
            all_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(
                    self.simplifier.simplified_and(
                        *(
                            self.simplifier.simplified_implies(guard, bound)
                            for bounds in all_bounds
                            for guard, bound in bounds
                        ),
                        transformed_body,
                    )
                )
            else:
                bodies.append(
                    self.simplifier.simplified_implies(
                        self.simplifier.simplified_and(
                            *(
                                self.simplifier.simplified_or(
                                    *(
                                        self.simplifier.simplified_and(
                                            guard,
                                            bound,
                                        )
                                        for guard, bound in element_bounds
                                    )
                                )
                                for element_bounds in all_bounds
                            )
                        ),
                        transformed_body,
                    )
                )

        base_variables: list[z3.Const] = [
            self._base_const(var) for var in bounding_vars
        ]

        if quantifier.is_exists():
            return z3.Exists(base_variables, self.simplifier.simplified_or(*bodies))
        else:
            return z3.ForAll(base_variables, self.simplifier.simplified_and(*bodies))

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

    def _get_bounded_elements(
        self, variables: list[z3.Const], elements: tuple[SymElement, ...]
    ) -> tuple[list[list[tuple[z3.BoolRef, z3.BoolRef]]], list[SymElement]]:
        all_bounds = []
        bounded_elements = []

        for variable, element in zip(variables, elements):
            base_variable = self._base_const(variable)
            if element.is_regular():
                bounds = [(self.true, base_variable == self.theory.regular_index())]
            elif element.sort.eq(z3.IntSort()):
                bounds = [(self.true, self.true)]
                element = element.with_index(base_variable)
            else:
                bounds = [
                    (
                        self.guards[element.node, guard].to_bool(),
                        bound,
                    )
                    for guard, bound in self.template.apply_bounds(
                        base_variable, scope=self.variables.scope(element.node)
                    ).items()
                ]
                assert bounds, f"No bounds for element {element}"
                element = element.with_index(base_variable)

            bounded_elements.append(element)
            all_bounds.append(bounds)

        return all_bounds, bounded_elements

    def _generate_aux(self, _default_assignment: Assignment) -> z3.BoolRef:
        with z3_as_z3():
            well_defined = well_definedness_formula(
                uninterpreted_constants(self.formula),
                uninterpreted_functions(self.formula),
            )

        well_defined_aux = self._transform_formula(
            well_defined, self.default_assignment
        )

        guards_aux = self._generate_guards_aux()

        return self.simplifier.simplified_and(well_defined_aux, guards_aux)

    def _generate_guards_aux(self) -> z3.BoolRef:
        return self.simplifier.simplified_and(
            *(
                exactly_one(
                    *(guard.to_bool() for guard in guards),
                )
                for _scope, guards in self.guards.items()
            )
        )

    def _get_template(self, relation: z3.FuncDeclRef | None) -> Template:
        if relation is None or all(
            not relation.domain(i).eq(z3.IntSort()) for i in range(relation.arity())
        ):
            return self.template
        else:
            assert config.fest.experimental.int_support(
                bool
            ), "`fest.experimental.int_support` is not enabled"
            if config.fest.experimental.regular_ints:
                return self.template
            return SPECIAL_FORM_INT


def _transpose_option(
    *pairs: tuple[Clause, SymElement]
) -> tuple[Clause | None, list[SymElement]]:
    clauses, elements = zip(*pairs)
    clause = reduce(operator.and_, clauses)
    return clause, list(elements)


@config(
    backend=config.fest.backend[str],
)
def find_sym_model(
    formula: z3.BoolRef,
    template: Template,
    *,
    backend: Backend = "auto",
    dry_run: bool = False,
) -> SymStruct | None:
    theory = template.theory_name

    base_theory_model: TheoryModel
    transformer: ModelFinder
    if backend == "cvc5" or (backend == "auto" and theory == "lia"):
        if config.fest.experimental.post_transformation_conversion(bool):
            assert not is_z3_as_cvc5()
            z3_transformer = ModelFinder(formula, template)
            base_theory_query = z3_transformer.result

            cvc5_lia_query = convert_z3_to_cvc5(base_theory_query)
            transformer = z3_transformer
        else:
            with z3_as_cvc5():
                cvc5_transformer = ModelFinder(formula, template)
                cvc5_lia_query = cast(cvc5.BoolRef, cvc5_transformer.result)

                transformer = cvc5_transformer
        cvc5_solver = cvc5_solver_for(transformer.theory.logic())
        print(
            cvc5_lia_query, file=(io.StringIO())
        )  # workaround for cvc5's Fatal failures
        check = cvc5_solver.check(cvc5_lia_query)
        if check != cvc5.sat:
            return None

        base_theory_model = transformer.theory.get_model(cvc5_solver.model())
    elif backend == "mona" or (backend == "auto" and theory == "string"):
        assert not is_z3_as_cvc5()
        z3_transformer = ModelFinder(formula, template)
        base_theory_query = z3_transformer.result
        transformer = z3_transformer

        mona_solver = PyMONAZ3Solver()
        if mona_solver.check(base_theory_query) != z3.sat:
            return None
        base_theory_model = transformer.theory.get_model(mona_solver.model())
    else:
        z3_transformer = ModelFinder(formula, template)

        base_theory_query = z3_transformer.result

        if dry_run:
            return None

        z3_solver = z3.SolverFor(z3_transformer.theory.logic())
        if z3_solver.check(base_theory_query) != z3.sat:
            return None

        transformer = z3_transformer
        base_theory_model = transformer.theory.get_model(z3_solver.model())

    return Model(
        base_model=base_theory_model,
        guards=transformer.guards,
        variables=transformer.variables,
        assignment=transformer.default_assignment,
        domain=transformer.domain,
        templates=transformer._get_template,
        order=transformer.order,
        formula=transformer.formula,
        theory_name=transformer.theory_name,
    ).as_sym_struct
