import pickle
from dataclasses import dataclass
from pathlib import Path
from typing import Self, Unpack, cast

import z3

from common.z3_utils.pickling import PicklableExpr, PicklableFunction, PicklableSort
from fest.symstructs import Assignment, BoundsFormula
from fest.theories import BaseTheory, TheoryCache, TheoryName, get_theory

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


@dataclass(eq=True, frozen=True)
class PicklableSymElement:
    node: Node
    sort: PicklableSort
    theory_name: TheoryName
    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, element.theory_name, None)
        else:
            return cls(
                element.node,
                sort,
                element.theory_name,
                PicklableExpr.create(element.index),
            )

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


@dataclass
class PicklableBoundsFormula(TheoryCache):
    theory_name: TheoryName
    expr: PicklableExpr

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

    @classmethod
    def from_bounds_formula(cls, theory: BaseTheory, formula: BoundsFormula) -> Self:
        return cls(
            theory.name(),
            PicklableExpr.create(formula(theory.to_base_const("i"))),
        )

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

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

        return bounds_formula


@dataclass
class PicklableSymTerm(TheoryCache):
    theory_name: TheoryName
    expr: PicklableExpr
    arity: int

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

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

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

        return substituting_term


@dataclass
class PicklableSymFormula(TheoryCache):
    theory_name: TheoryName
    expr: PicklableExpr
    arity: int

    @classmethod
    def from_sym_formula(
        cls, theory: TheoryName | BaseTheory, formula: SymFormula, arity: int
    ) -> Self:
        if not isinstance(theory, BaseTheory):
            theory = get_theory(theory)
        indices = [theory.to_base_const(f"i_{i + 1}") for i in range(arity)]
        expr = PicklableExpr.create(formula(*indices))
        return cls(theory.name(), expr, arity)

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

        def substituting_formula(*indices: z3.ExprRef) -> 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]
    theory_name: TheoryName

    @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(struct.theory, formula)
                for node, formula in struct.bounds.items()
            },
            functions={
                (PicklableFunction.create(fun), *nodes): (
                    target,
                    PicklableSymTerm.from_sym_term(struct.theory, term, len(nodes)),
                )
                for (fun, *nodes), (target, term) in struct.functions.items()
            },
            relations={
                PicklableFunction.create(relation): {
                    nodes: PicklableSymFormula.from_sym_formula(
                        struct.theory, 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()
            },
            theory_name=struct.theory_name,
        )

    @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()
            },
            theory_name=self.theory_name,
        )


@dataclass
class PicklableAssignment:
    values: dict[str, PicklableSymElement]

    @classmethod
    def from_assignment(cls, assignment: Assignment) -> Self:
        return cls(
            {
                key: PicklableSymElement.from_element(element)
                for key, element in assignment.items()
            }
        )

    @property
    def assignment(self) -> Assignment:
        return Assignment(
            {key: element.element for key, element in self.values.items()}
        )


def save_struct(struct: SymStruct, path: Path) -> None:
    with path.open("wb") as file:
        pickle.dump(PicklableSymStruct.from_struct(struct), file)


def load_struct(path: Path) -> SymStruct:
    with path.open("rb") as file:
        return cast(PicklableSymStruct, pickle.load(file)).struct
