from dataclasses import dataclass, field, replace
from typing import Protocol, Self, cast, overload

import z3

from common.z3_utils.expr_utils import transform_quantifier

type QID = str
type Variable = str


class QIDInstantiationBuilder(Protocol):
    def __call__(self, **kwargs: z3.ExprRef | int) -> "Instantiation": ...


@dataclass
class Instantiation:
    mapping: dict[tuple[QID, Variable], z3.ExprRef] = field(default_factory=dict)
    drop_quantifiers: bool = False

    def instantiate(self, formula: z3.BoolRef) -> z3.BoolRef:
        if z3.is_quantifier(formula):
            qid = formula.qid()

            if qid == "<null>":
                # No qid, keep quantifier
                return transform_quantifier(
                    formula, lambda body: self.instantiate(body)
                )

            variables = [(qid, formula.var_name(i)) for i in range(formula.num_vars())]

            has_instantiations = list(map(self.mapping.__contains__, variables))

            if not any(has_instantiations):
                if self.drop_quantifiers:
                    if formula.is_forall():
                        return z3.BoolVal(True)
                    else:
                        return z3.BoolVal(False)
                # No instantiations, keep quantifier
                return transform_quantifier(
                    formula, lambda body: self.instantiate(body)
                )
            elif all(has_instantiations):
                substitutions = reversed(list(map(self.mapping.__getitem__, variables)))

                return self.instantiate(
                    z3.substitute_vars(formula.body(), *substitutions)
                )
            else:
                # partial instantiation
                missing_variables = ", ".join(
                    f"{qid}:{variable}"
                    for qid, variable in set(variables) - set(self.mapping.keys())
                )
                assert False, f"Partial instantiation, missing: {missing_variables}"

        elif (
            z3.is_app(formula)
            and not z3.is_const(formula)
            and z3.is_bool(formula.children()[0])
            and formula.decl().kind() != z3.Z3_OP_UNINTERPRETED
        ):
            instantiated_children = [
                self.instantiate(cast(z3.BoolRef, child))
                for child in formula.children()
            ]

            if z3.is_eq(formula):
                left, right = instantiated_children
                return left == right
            elif z3.is_distinct(formula):
                left, right = instantiated_children
                return left != right

            return cast(
                z3.BoolRef, getattr(z3, str(formula.decl()))(*instantiated_children)
            )
        else:
            return formula

    def __getattr__(self, item: str) -> QIDInstantiationBuilder:
        def put_instantiation(**kwargs: z3.ExprRef | int) -> Instantiation:
            for variable, value in kwargs.items():
                if isinstance(value, int):
                    value = z3.IntVal(value)
                self.mapping[(item, variable)] = value

            return self

        return put_instantiation

    @overload
    def __and__(self, other: Self) -> Self: ...

    @overload
    def __and__(self, other: "InstantiationSet") -> "InstantiationSet": ...

    def __and__(self, other: Self | "InstantiationSet") -> Self | "InstantiationSet":
        if isinstance(other, InstantiationSet):
            return other & InstantiationSet([self])
        return replace(self, mapping=self.mapping | other.mapping)

    def __or__(self, other: Self) -> "InstantiationSet":
        return InstantiationSet([self, other])


@dataclass
class InstantiationSet:
    instantiations: list[Instantiation] = field(default_factory=list)

    def __and__(self, other: Self) -> Self:
        return replace(
            self,
            instantiations=[
                i1 & i2 for i1 in self.instantiations for i2 in other.instantiations
            ],
        )

    @overload
    def __or__(self, other: Instantiation) -> Self: ...
    @overload
    def __or__(self, other: Self) -> Self: ...
    def __or__(self, other: Instantiation | Self) -> Self:
        if isinstance(other, Instantiation):
            return replace(self, instantiations=self.instantiations + [other])
        else:
            return replace(
                self, instantiations=self.instantiations + other.instantiations
            )

    def instantiate(self, formula: z3.BoolRef) -> list[z3.BoolRef]:
        return [i.instantiate(formula) for i in self.instantiations]


def assert_qids(formula: z3.BoolRef) -> None:
    seen_qids = set()

    def recursive_assert_qids(sub_formula: z3.BoolRef) -> None:
        if z3.is_quantifier(sub_formula):
            qid = sub_formula.qid()
            assert qid != "<null>", f"Missing qid for {sub_formula}"
            assert qid not in seen_qids, f"Duplicate qid `{qid}`"
            seen_qids.add(qid)
            recursive_assert_qids(sub_formula.body())
        elif (
            z3.is_app(sub_formula)
            and z3.is_bool(sub_formula.children()[0])
            and sub_formula.decl().kind() != z3.Z3_OP_UNINTERPRETED
        ):
            for child in sub_formula.children():
                recursive_assert_qids(cast(z3.BoolRef, child))

    recursive_assert_qids(formula)
