from functools import cached_property
from typing import Literal, Protocol, runtime_checkable

import z3

from . import arith, string


class TheoryModel(Protocol):
    def eval_bool(self, item: str) -> bool: ...
    def eval_expr(self, item: str) -> z3.ExprRef: ...


type TheoryName = Literal["lia", "string"]


@runtime_checkable
class BaseTheory(Protocol):
    def name(self) -> TheoryName: ...

    def regular_index(self) -> z3.ExprRef: ...

    def to_base_const(self, const: str) -> z3.Const: ...

    def logic(self) -> str: ...

    def get_model[ModelRef](self, model: ModelRef) -> TheoryModel: ...


def get_theory(name: TheoryName) -> BaseTheory:
    if name == "lia":
        return arith
    elif name == "string":
        return string
    raise ValueError(f"Unsupported theory `{name}`")


class TheoryNamed(Protocol):
    theory_name: TheoryName


class TheoryCache(TheoryNamed):
    @cached_property
    def theory(self) -> BaseTheory:
        return get_theory(self.theory_name)


class TheoryCacheDelegate(Protocol):
    @property
    def delegate(self) -> TheoryNamed: ...

    @cached_property
    def theory_name(self) -> TheoryName:
        return self.delegate.theory_name

    @cached_property
    def theory(self) -> BaseTheory:
        return get_theory(self.theory_name)
