from dataclasses import dataclass
from typing import Self

import z3

from common.relations import Relation, is_relation


@dataclass(eq=True, frozen=True, slots=True)
class PicklableExpr:
    sexpr: str

    @classmethod
    def create(cls, expr: z3.ExprRef) -> Self:
        try:
            solver = z3.Solver()
            solver.add(expr == expr)
            return cls(solver.sexpr())
        except Exception:
            raise RuntimeError(f"Could not pickle: {expr.__class__}")

    @property
    def expr(self) -> z3.ExprRef:
        try:
            solver = z3.Solver()
            solver.from_string(self.sexpr)
            return solver.assertions()[0].children()[0]
        except Exception:
            raise RuntimeError(f"Could not parse {self.sexpr}")

    @property
    def formula(self) -> z3.BoolRef:
        assert z3.is_bool(self.expr)
        return self.expr

    @property
    def const(self) -> z3.Const:
        assert z3.is_const(self.expr)
        return self.expr

    @property
    def arith(self) -> z3.ArithRef:
        assert z3.is_arith(self.expr)
        return self.expr


@dataclass(eq=True, frozen=True, slots=True)
class PicklableSort:
    sort_name: str

    @classmethod
    def create(cls, sort: z3.SortRef) -> Self:
        return cls(sort.name())

    @property
    def sort(self) -> z3.SortRef:
        if self.sort_name == "Int":
            return z3.IntSort()
        elif self.sort_name == "Bool":
            return z3.BoolSort()
        return z3.DeclareSort(self.sort_name)


@dataclass(eq=True, frozen=True, slots=True)
class PicklableFunction:
    fun_name: str
    signature: tuple[PicklableSort, ...]

    @classmethod
    def create(cls, fun: z3.FuncDeclRef) -> Self:
        name = fun.name()
        signature = tuple(
            PicklableSort.create(fun.domain(i)) for i in range(fun.arity())
        ) + (PicklableSort.create(fun.range()),)
        return cls(name, signature)

    @property
    def fun(self) -> z3.FuncDeclRef:
        return z3.Function(self.fun_name, *[sort.sort for sort in self.signature])

    @property
    def relation(self) -> Relation:
        assert is_relation(self.fun)
        return self.fun
