from z3 import *
from pydantic import BaseModel
from pydantic.config import ConfigDict
from typing import Type, TypedDict, Never


class ArbTypesModel(BaseModel):
    model_config: ConfigDict = ConfigDict(arbitrary_types_allowed=True)


def free_vars(expr):
    s = set()
    stack = [expr]
    while stack:
        e = stack.pop()
        if is_const(e) and e.decl().kind() == Z3_OP_UNINTERPRETED:
            s.add(e)
        for c in e.children():
            stack.append(c)
    return s


class SimpleDatatypes(TypedDict):
    int: ArithSortRef
    float: FPSortRef
    bool: BoolSortRef
    str: SeqSortRef


def get_simple_datatypes() -> SimpleDatatypes:
    return {
        int: IntSort(),
        float: Float64(),
        bool: BoolSort(),
        str: StringSort(),
    }


class OptDatatypes(TypedDict):
    int: DatatypeSortRef
    float: DatatypeSortRef
    bool: DatatypeSortRef
    str: DatatypeSortRef


def get_optional_datatypes(simple_types: SimpleDatatypes) -> OptDatatypes:

    def _build_opt_dt(raw_sort: z3.SortRef, typename: str) -> DatatypeSortRef:
        NewType = Datatype('Opt' + typename)
        NewType.declare('Nil')
        NewType.declare('Some', ('val', raw_sort))
        return NewType.create()

    return {
        py_type: _build_opt_dt(z3_raw_sort, py_type.__name__.capitalize())
        for py_type, z3_raw_sort in simple_types.items()
    }


SimpleTypesZ3 = get_simple_datatypes()
OptTypesZ3 = get_optional_datatypes(SimpleTypesZ3)
