import io
from collections import defaultdict
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from functools import cached_property
from itertools import product
from typing import Iterable, Unpack

import z3

from common import config
from common.relations import Relation
from common.z3_utils.expr_utils import (
    constants,
    uninterpreted_functions,
    uninterpreted_relations,
)
from fest import DUMMY_SORT
from fest.symstructs import (
    Assignment,
    Node,
    SymElement,
    SymFormula,
    SymStruct,
    SymTerm,
    default_bounds,
)
from fest.theories import TheoryCache, TheoryModel, TheoryName

from . import Template
from .guards import Guard
from .template_free_variables import TemplateFreeVariables
from .utils import lexicographical


@dataclass(slots=True)
class TemplateModelFreeVariables(TemplateFreeVariables):
    theory_model: TheoryModel
    theory_name: TheoryName

    def __getitem__(self, item: tuple[str, str]) -> z3.ExprRef:
        prefix, suffix = item
        return self.theory_model.eval_expr(prefix + suffix)


@dataclass(kw_only=True)
class Model(TheoryCache):
    base_model: TheoryModel
    guards: Mapping[str, set[Guard]]
    variables: Iterable[str]
    domain: dict[z3.SortRef, set[SymElement]]
    assignment: Assignment
    templates: Callable[[Relation | None], Template]
    order: Relation | None
    formula: z3.BoolRef
    theory_name: TheoryName

    @cached_property
    def template_free_variables(self) -> TemplateFreeVariables:
        return TemplateModelFreeVariables(self.base_model, self.theory_name)

    @cached_property
    def all_elements(self) -> set[SymElement]:
        return {element for elements in self.domain.values() for element in elements}

    @cached_property
    def relations(self) -> set[Relation]:
        return uninterpreted_relations(self.formula)

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

    @cached_property
    def constants(self) -> frozenset[z3.Const]:
        return constants(self.formula)

    @cached_property
    def as_sym_struct(self) -> SymStruct:
        domain = {
            sort: {element.node for element in elements}
            for sort, elements in self.domain.items()
        }

        bounds = default_bounds()
        for element in self.all_elements:
            if element.is_regular() or element.sort.eq(z3.IntSort()):
                continue
            node = element.node
            guard = self._eval_scope(node)
            bounds[element.node] = self.templates(None)[
                "bound",
                guard.name,
                self.template_free_variables.scope(node),
            ]

        consts = {}
        for constant in self.constants:
            consts[constant] = self._get_constant(constant)

        functions = {}
        for function in self.functions:
            arg_domains = [
                self.domain[function.domain(i)] for i in range(function.arity())
            ]
            for args in product(*arg_domains):
                target = self._get_target(function, *args)

                if self._get_element(function.range(), target).is_regular():
                    sym_term: SymTerm = lambda *indices: self.theory.regular_index()
                else:
                    sym_term = self._stretch_term(
                        self._get_sym_term(function, *args), *args
                    )

                nodes = tuple(arg.node for arg in args)
                function_key: tuple[z3.FuncDeclRef, Unpack[tuple[Node, ...]]] = (
                    function,
                    *nodes,
                )

                functions[function_key] = (target, sym_term)

        relations: dict[Relation, dict[tuple[Node, ...], SymFormula]] = defaultdict(
            dict
        )
        for relation in self.relations:
            arg_domains = [
                self.domain[relation.domain(i)] for i in range(relation.arity())
            ]
            for args in product(*arg_domains):
                formula = self._get_sym_formula(relation, *args)

                nodes = tuple(arg.node for arg in args)
                relations[relation][nodes] = self._stretch_formula(formula, *args)

        return SymStruct(
            domain=domain,
            bounds=bounds,
            functions=functions,
            relations=relations,
            constants=consts,
            theory_name=self.theory_name,
        )

    def __str__(self) -> str:
        buffer = io.StringIO()
        for sort, elements in self.domain.items():
            if sort.eq(z3.IntSort()) or sort.name() == DUMMY_SORT:
                continue

            print(f"{sort}: {", ".join(map(str, elements))}", file=buffer)

        for const, element in self.assignment.items():
            print(f"{const}:{element}", file=buffer)

        for scope in self.guards:
            print(str(self._eval_scope(scope)), file=buffer)

        for variable in self.variables:
            varname = self.theory.to_base_const(variable).decl().name()
            print(f"{variable} = {self.base_model.eval_expr(varname)}", file=buffer)

        return buffer.getvalue()

    def _eval_scope(self, scope: str) -> Guard:
        chosen = [
            guard
            for guard in self.guards[scope]
            if self.base_model.eval_bool(str(guard))
        ]

        assert (
            len(chosen) == 1
        ), f"Inconsistent model: scope {scope} has guards {chosen} out of {self.guards[scope]}"

        return chosen[0]

    def _get_sym_formula(self, relation: Relation, *elements: SymElement) -> SymFormula:
        if self.order is not None and self.order.eq(relation):
            assert len(elements) == 2, "Non-binary order evaluation"
            return self._get_order(*elements)
        else:
            compact_arity = len(
                [element for element in elements if element.is_summary()]
            )
            nodes = ", ".join(element.node for element in elements)
            scope = f"{relation}({nodes})"
            guard = self._eval_scope(scope)
            return self.templates(relation)[
                "rel",
                guard.name,
                compact_arity,
                self.template_free_variables.scope(scope),
            ]

    def _get_order(self, left: SymElement, right: SymElement) -> SymFormula:
        is_lexicographical, first, second = lexicographical(left.sort, left, right)
        nodes = f"{first.node}, {second.node}"

        scope = f"{self.order}({nodes})"
        guard = self._eval_scope(scope)

        match guard.name, is_lexicographical:
            case ("top", True) | ("bottom", False):
                return lambda *indices: z3.BoolVal(True)
            case ("top", False) | ("bottom", True):
                return lambda *indices: z3.BoolVal(False)
            case ("lt", True) | ("leq", False):
                return lambda *indices: indices[0] < indices[1]
            case ("lt", False) | ("leq", True):
                return lambda *indices: indices[0] <= indices[1]
            case _:
                assert False, f"Unexpected order formula {guard.name}"

    def _get_target(self, function: z3.FuncDeclRef, *args: SymElement) -> Node:
        return self._eval_scope(
            f"{function}({", ".join(arg.node for arg in args)})->"
        ).name

    def _get_sym_term(self, function: z3.FuncDeclRef, *elements: SymElement) -> SymTerm:
        compact_arity = len([element for element in elements if element.is_summary()])

        nodes = ", ".join(element.node for element in elements)
        scope = f"{function}({nodes})"
        guard = self._eval_scope(scope)
        return self.templates(None)[
            "fun",
            guard.name,
            compact_arity,
            self.template_free_variables.scope(scope),
        ]

    @staticmethod
    def _stretch_formula(formula: SymFormula, *args: SymElement) -> SymFormula:
        # As an optimization, templates use only summary nodes' indices
        # in sym-formulas.
        # Therefore, we need to stretch out the indices

        summary_indices = [i for i, arg in enumerate(args) if arg.is_summary()]

        def stretched_formula(*indices: z3.ExprRef) -> z3.BoolRef:
            used_indices = [indices[i] for i in summary_indices]
            return formula(*used_indices)

        return stretched_formula

    @staticmethod
    def _stretch_term(term: SymTerm, *args: SymElement) -> SymTerm:
        # As an optimization, templates use only summary nodes' indices
        # in sym-terms.
        # Therefore, we need to stretch out the indices

        summary_indices = [i for i, arg in enumerate(args) if arg.is_summary()]

        def stretched_term(*indices: z3.ExprRef) -> z3.ExprRef:
            used_indices = [indices[i] for i in summary_indices]
            return term(*used_indices)

        return stretched_term

    def _get_constant(self, constant: z3.Const) -> SymElement:
        name = str(constant)
        if constant.sort().eq(z3.IntSort()):
            assert config.fest.experimental.int_support(
                bool
            ), "`fest.experimental.int_support` is not enabled"
            value = self.base_model.eval_expr(name)
            assert z3.is_arith(value)
            return SymElement.int(value)
        if name in self.assignment:
            return self.assignment[name]
        else:
            element = self._get_element(constant.sort(), self._eval_scope(name).name)
            assert (
                element is not None
            ), f"Inconsistent model, no element for constant `{constant}`"

            if element.is_regular():
                return element

            index = self.template_free_variables[
                constant.decl().name(),
                "i",
            ]
            return element.with_index(index)

    def _get_element(self, sort: z3.SortRef, node: str) -> SymElement:
        element = next(
            (element for element in self.domain[sort] if element.node == node),
            None,
        )
        assert element is not None, f"Missing element for `{node}` of sort `{sort}`"
        return element
