from dataclasses import dataclass
from typing import Self, Unpack

import z3

from common.z3_utils.pickling import PicklableExpr, PicklableFunction, PicklableSort

from . import BoundsFormula, Node, SymElement, SymFormula, SymStruct, SymTerm


@dataclass(eq=True, frozen=True)
class PicklableSymElement:
    node: Node
    sort: PicklableSort
    index: PicklableExpr | None

    @classmethod
    def from_element(cls, element: SymElement) -> Self:
        sort = PicklableSort.create(element.sort)
        if element.index is None:
            return cls(element.node, sort, None)
        else:
            return cls(element.node, sort, PicklableExpr.create(element.index))

    @property
    def element(self) -> SymElement:
        sort = self.sort.sort
        if self.index is None:
            return SymElement(self.node, sort)
        else:
            return SymElement(self.node, sort, self.index.arity)


@dataclass
class PicklableBoundsFormula:
    expr: PicklableExpr

    def __call__(self, index: z3.ArithRef) -> z3.BoolRef:
        return self.formula(index)

    @classmethod
    def from_bounds_formula(cls, formula: BoundsFormula) -> Self:
        return cls(PicklableExpr.create(formula(z3.Int("i"))))

    @property
    def formula(self) -> BoundsFormula:
        expr = self.expr.expr
        assert z3.is_bool(expr)
        i = z3.Int("i")

        def bounds_formula(index: z3.ArithRef) -> z3.BoolRef:
            return z3.substitute(expr, (i, index))

        return bounds_formula


@dataclass
class PicklableSymTerm:
    expr: PicklableExpr
    arity: int

    @classmethod
    def from_sym_term(cls, term: SymTerm, arity: int) -> Self:
        indices = [z3.Int(f"i_{i + 1}") for i in range(arity)]
        expr = PicklableExpr.create(term(*indices))
        return cls(expr, arity)

    @property
    def sym_term(self) -> SymTerm:
        default_indices = [z3.Int(f"i_{i + 1}") for i in range(self.arity)]
        expr = self.expr.arity

        def substituting_term(*indices: z3.ArithRef) -> z3.ArithRef:
            return z3.substitute(expr, *zip(default_indices, indices))

        return substituting_term


@dataclass
class PicklableSymFormula:
    expr: PicklableExpr
    arity: int

    @classmethod
    def from_sym_formula(cls, formula: SymFormula, arity: int) -> Self:
        indices = [z3.Int(f"i_{i + 1}") for i in range(arity)]
        expr = PicklableExpr.create(formula(*indices))
        return cls(expr, arity)

    @property
    def sym_formula(self) -> SymFormula:
        formula = self.expr.formula
        arity = self.arity
        default_indices = [z3.Int(f"i_{i + 1}") for i in range(arity)]

        def substituting_formula(*indices: z3.ArithRef) -> z3.BoolRef:
            return z3.substitute(formula, *zip(default_indices, indices))

        return substituting_formula


@dataclass(kw_only=True)
class PicklableSymStruct:
    domain: dict[PicklableSort, set[Node]]
    bounds: dict[Node, PicklableBoundsFormula]
    functions: dict[
        tuple[PicklableFunction, Unpack[tuple[Node, ...]]],
        tuple[Node, PicklableSymTerm],
    ]
    relations: dict[PicklableFunction, dict[tuple[Node, ...], PicklableSymFormula]]
    constants: dict[PicklableExpr, PicklableSymElement]

    @classmethod
    def from_struct(cls, struct: SymStruct) -> Self:
        return cls(
            domain={
                PicklableSort.create(sort): nodes
                for sort, nodes in struct.domain.items()
            },
            bounds={
                node: PicklableBoundsFormula.from_bounds_formula(formula)
                for node, formula in struct.bounds.items()
            },
            functions={
                (PicklableFunction.create(fun), *nodes): (
                    target,
                    PicklableSymTerm.from_sym_term(term, len(nodes)),
                )
                for (fun, *nodes), (target, term) in struct.functions.items()
            },
            relations={
                PicklableFunction.create(relation): {
                    nodes: PicklableSymFormula.from_sym_formula(definition, len(nodes))
                    for nodes, definition in definitions.items()
                }
                for relation, definitions in struct.relations.items()
            },
            constants={
                PicklableExpr.create(const): PicklableSymElement.from_element(element)
                for const, element in struct.constants.items()
            },
        )

    @property
    def struct(self) -> SymStruct:
        return SymStruct(
            domain={sort.sort: nodes for sort, nodes in self.domain.items()},
            bounds={node: formula.formula for node, formula in self.bounds.items()},
            functions={
                (fun.fun, *nodes): (target, term.sym_term)
                for (fun, *nodes), (target, term) in self.functions.items()
            },
            relations={
                relation.relation: {
                    nodes: definition.sym_formula
                    for nodes, definition in definitions.items()
                }
                for relation, definitions in self.relations.items()
            },
            constants={
                const.const: element.element
                for const, element in self.constants.items()
            },
        )
