from dataclasses import dataclass
from typing import Literal

import z3
from cvc5 import pythonic as cvc5


@dataclass(slots=True)
class Z3LiaModel:
    z3_model: z3.ModelRef

    def eval_bool(self, item: str) -> bool:
        guard_result = self.z3_model.eval(z3.Bool(item), model_completion=True)

        if z3.is_true(guard_result):
            return True
        elif z3.is_false(guard_result):
            return False

        assert False, f"Incomplete model: `{item}` missing"

    def eval_expr(self, item: str) -> z3.ArithRef:
        variable_result = self.z3_model.eval(z3.Int(item), model_completion=True)

        assert z3.is_int_value(
            variable_result
        ), f"Incomplete model: `{item} = {variable_result}`"

        return variable_result


@dataclass(slots=True)
class Cvc5LiaModel:
    cvc5_model: cvc5.ModelRef

    def eval_bool(self, item: str) -> bool:
        guard_result = self.cvc5_model.eval(cvc5.Bool(item), model_completion=True)

        if cvc5.is_true(guard_result):
            return True
        elif cvc5.is_false(guard_result):
            return False
        raise ValueError(f"Incomplete model: missing `{item}`")

    def eval_expr(self, item: str) -> z3.ArithRef:
        variable_result = self.cvc5_model.eval(cvc5.Int(item), model_completion=True)

        assert cvc5.is_int_value(
            variable_result
        ), f"Incomplete model: `{item} = {variable_result}`"
        return z3.IntVal(variable_result.as_long())


def name() -> Literal["lia"]:
    return "lia"


def regular_index() -> z3.ArithRef:
    return z3.IntVal(0)


def to_base_const(const: str) -> z3.Int:
    return z3.Int(str(const))


def logic() -> str:
    return "LIA"


def get_model[ModelRef](model: ModelRef) -> Z3LiaModel | Cvc5LiaModel:
    if isinstance(model, z3.ModelRef):
        return Z3LiaModel(model)
    elif isinstance(model, cvc5.ModelRef):
        return Cvc5LiaModel(model)
    raise ValueError(f"Unsupported model `{model.__class__.__name__}`")
