from collections.abc import Iterator
from contextlib import contextmanager
from typing import Any

import z3
from cvc5 import pythonic as cvc5

_CONSTRUCTORS_TO_PATCH = [
    "Const",
    "Consts",
    "Int",
    "Ints",
    "IntVal",
    "String",
    "Strings",
    "StringVal",
    "Bool",
    "Bools",
    "BoolVal",
    "And",
    "Or",
    "Implies",
    "Not",
    "Xor",
    "Exists",
    "ForAll",
    "If",
    "PrefixOf",
    "Concat",
]


_PATCHED_MARKER = "_patched"

_UNPATCHED_Z3: dict[str, Any] = {}

_unique = object()


@contextmanager
def z3_as_cvc5() -> Iterator[None]:
    _patch_z3_to_cvc5()
    try:
        yield
    finally:
        _restore_z3()


@contextmanager
def z3_as_z3() -> Iterator[None]:
    if is_z3_as_cvc5():
        _restore_z3()
        try:
            yield
        finally:
            _patch_z3_to_cvc5()
    else:
        yield


def is_z3_as_cvc5() -> bool:
    return hasattr(z3, _PATCHED_MARKER) and getattr(z3, _PATCHED_MARKER) is _unique


def _patch_z3_to_cvc5() -> None:
    if is_z3_as_cvc5():
        return

    for attr in _CONSTRUCTORS_TO_PATCH:
        _patch_attr(attr)

    setattr(z3, _PATCHED_MARKER, _unique)


def _restore_z3() -> None:
    if not is_z3_as_cvc5():
        return

    for attr in _CONSTRUCTORS_TO_PATCH:
        _restore_attr(attr)

    delattr(z3, _PATCHED_MARKER)


def _patch_attr(attr: str) -> None:
    assert hasattr(z3, attr), f"Could not patch missing attr `{attr}`"
    assert hasattr(cvc5, attr), f"Could not find analog attr `{attr}`"
    original = getattr(z3, attr)
    cvc5_analog = getattr(cvc5, attr)

    if attr not in _UNPATCHED_Z3:
        _UNPATCHED_Z3[attr] = original
    setattr(z3, attr, cvc5_analog)


def _restore_attr(attr: str) -> None:
    assert (
        hasattr(z3, attr) and attr in _UNPATCHED_Z3
    ), f"Could not unpatch missing attr `{attr}`"
    setattr(z3, attr, _UNPATCHED_Z3[attr])


def polymorphic_check(attr: str, value: Any) -> bool:
    assert hasattr(z3, attr), f"Could not find tester `{attr}`"
    assert hasattr(cvc5, attr), f"Could not find analog tester `{attr}`"

    z3_tester = getattr(z3, attr)
    cvc5_tester = getattr(cvc5, attr)

    if not hasattr(value, "__module__"):
        return False
    if value.__module__.startswith("z3."):
        return z3_tester(value) is True
    else:
        return cvc5_tester(value) is True
