# Copyright (C) 2022-2024 Exaloop Inc. <https://exaloop.io>

@tuple
@__internal__
class type[T]:
    #  __new__ / __init__: always done internally
    # __repr__
    # __call__
    # __name__
    # __module__
    # __doc__
    pass

@tuple
@__internal__
class unrealized_type[T]:
    pass

@__internal__
class __internal__:
    pass
@__internal__
class __magic__:
    pass

@tuple
@__internal__
@__notuple__
class bool:
    pass

@tuple
@__internal__
@__notuple__
class byte:
    pass

@tuple
@__internal__
@__notuple__
class int:
    MAX = 9223372036854775807
    pass

@tuple
@__internal__
@__notuple__
class float:
    MIN_10_EXP = -307
    pass

@tuple
@__internal__
@__notuple__
class float32:
    MIN_10_EXP = -37
    pass

@tuple
@__internal__
@__notuple__
class float16:
    MIN_10_EXP = -4
    pass

@tuple
@__internal__
@__notuple__
class bfloat16:
    MIN_10_EXP = -37
    pass

@tuple
@__internal__
@__notuple__
class float128:
    MIN_10_EXP = -4931
    pass

@tuple
@__internal__
@__notuple__
class Function[T, TR]:
    pass

@tuple
@__internal__
class Callable[T, TR]:
    pass

@tuple
@__internal__
class NoneType:
    pass

@tuple
@__internal__
@__notuple__
class Ptr[T]:
    pass
cobj = Ptr[byte]

@tuple
@__internal__
@__notuple__
class Generator[T]:
    pass

@tuple
@__internal__
@__notuple__
class Optional:
    T: type = NoneType

@tuple
@__internal__
@__notuple__
class Int[N: Static[int]]:
    pass

@tuple
@__internal__
@__notuple__
class UInt[N: Static[int]]:
    pass

@__internal__
class pyobj:
    p: Ptr[byte]

@tuple
@__internal__
class str:
    ptr: Ptr[byte]
    len: int

@tuple
@__internal__
class Tuple:
    def __add__(self: __SELF__, obj, __SELF__: type):
        return __magic__.add(self, obj)
    def __mul__(self: __SELF__, n: Static[int], __SELF__: type):
        return __magic__.mul(self, n)
    def __contains__(self: __SELF__, obj, __SELF__: type) -> bool:
        return __magic__.contains(self, obj)
    def __getitem__(self: __SELF__, idx: int, __SELF__: type):
        return __magic__.getitem(self, idx)
    def __iter__(self: __SELF__, __SELF__: type):
        yield from __magic__.iter(self)
    def __hash__(self: __SELF__, __SELF__: type) -> int:
        return __magic__.hash(self)
    def __repr__(self: __SELF__, __SELF__: type) -> str:
        return __magic__.repr(self)
    def __len__(self: __SELF__, __SELF__: type) -> int:
        return __magic__.len(self)
    def __eq__(self: __SELF__, obj: __SELF__, __SELF__: type) -> bool:
        return __magic__.eq(self, obj)
    def __ne__(self: __SELF__, obj: __SELF__, __SELF__: type) -> bool:
        return __magic__.ne(self, obj)
    def __gt__(self: __SELF__, obj: __SELF__, __SELF__: type) -> bool:
        return __magic__.gt(self, obj)
    def __ge__(self: __SELF__, obj: __SELF__, __SELF__: type) -> bool:
        return __magic__.ge(self, obj)
    def __lt__(self: __SELF__, obj: __SELF__, __SELF__: type) -> bool:
        return __magic__.lt(self, obj)
    def __le__(self: __SELF__, obj: __SELF__, __SELF__: type) -> bool:
        return __magic__.le(self, obj)
    def __pickle__(self: __SELF__, dest: Ptr[byte], __SELF__: type):
        return __magic__.pickle(self, dest)
    def __unpickle__(src: Ptr[byte], __SELF__: type) -> __SELF__:
        return __magic__.unpickle(src)
    def __to_py__(self: __SELF__, __SELF__: type) -> Ptr[byte]:
        return __magic__.to_py(self)
    def __from_py__(src: Ptr[byte], __SELF__: type) -> __SELF__:
        return __magic__.from_py(src)
    def __to_gpu__(self: __SELF__, cache, __SELF__: type) -> __SELF__:
        return __magic__.to_gpu(self, cache)
    def __from_gpu__(self: __SELF__, other: __SELF__, __SELF__: type):
        return __magic__.from_gpu(self, other)
    def __from_gpu_new__(other: __SELF__, __SELF__: type) -> __SELF__:
        return __magic__.from_gpu_new(other)
    def __tuplesize__(self: __SELF__, __SELF__: type) -> int:
        return __magic__.tuplesize(self)
tuple = Tuple

@__attribute__
def pure():
    pass

@__attribute__
def derives():
    pass

@extend
class NoneType:
    @pure
    @derives
    @llvm
    def __new__() -> NoneType:
        ret {} {}

@tuple
@__internal__
class Array:
    len: int
    ptr: Ptr[T]
    T: type

@extend
class type:
    def __new__(obj):
        pass
function = Function

@__internal__
class Ref[T]:
    pass

@tuple
@__internal__
@__notuple__
class Union[TU]:
    # compiler-generated
    def __new__(val):
        TU
    def __call__(self, *args, **kwargs):
        return __internal__.union_call(self, args, kwargs)

@extend
class Function:
    @pure
    @derives
    @llvm
    def __new__() -> Function[T, TR]:
        ret ptr null

# dummy
@__internal__
class TypeVar[T]: pass
@__internal__
class ByVal: pass
@__internal__
class ByRef: pass

@__internal__
class ClassVar[T]:
    pass

@__internal__
class RTTI:
    id: int

@__internal__
@tuple
class ellipsis:
    @pure
    @derives
    @llvm
    def __new__() -> ellipsis:
        ret {} {}
Ellipsis = ellipsis()

@tuple
@__internal__
class __array__:
    T: type
    def __new__(sz: Static[int]) -> Array[T]:
        pass

@tuple
@__internal__
class Import:
    __ipath__: str
    __iname__: str
    P: Static[str]

    @pure
    @derives
    @llvm
    def __new__(path: str, name: str, P: Static[str]) -> Import[P]:
        %0 = insertvalue { {=str}, {=str} } undef, {=str} %path, 0
        %1 = insertvalue { {=str}, {=str} } %0,    {=str} %name, 1
        ret { {=str}, {=str} } %1

def __ptr__(var):
    pass

def staticlen(obj):
    pass

def compile_error(msg: Static[str]):
    pass

def isinstance(obj, what):
    pass

@__attribute__
def overload():
    pass

def hasattr(obj, attr: Static[str], *args, **kwargs):
    """Special handling"""
    pass

def getattr(obj, attr: Static[str]):
    pass

def setattr(obj, attr: Static[str], what):
    pass

def super():
    pass

def superf(*args):
    """Special handling"""
    pass

def __realized__(fn, args):
    pass

def statictuple(*args):
    return args

def __has_rtti__(T: type):
    pass

#(init=False, repr=False, eq=False, order=False, hash=False, pickle=False, python=False, gpu=False, container=False)
@__internal__
@tuple
class NamedTuple:
    args: T
    N: Static[int]  # name cache ID
    T: type

    @pure
    @derives
    @llvm
    def __new__(args: T = (), N: Static[int] = 0, T: type) -> NamedTuple[N, T]:
        %0 = insertvalue { {=T} } undef, {=T} %args, 0
        ret { {=T} } %0

    def __getitem__(self, key: Static[str]):
        return getattr(self, key)

    def __contains__(self, key: Static[str]):
        return hasattr(self, key)

    def get(self, key: Static[str], default = NoneType()):
        return __internal__.kwargs_get(self, key, default)

    def __keys__(self):
        return __internal__.namedkeys(N)

    def __repr__(self):
        keys = self.__keys__()
        values = [v.__repr__() for v in self.args]
        s = ', '.join(f"{keys[i]}: {values[i]}" for i in range(len(keys)))
        return f"({s})"

@__internal__
@tuple
class Partial:
    args: T  # format: (arg1, arg2, ..., (star_args...))
    kwargs: K

    M: Static[str] # mask
    T: type # Tuple
    K: type # NamedTuple
    F: type # must be unrealized_type

    @pure
    @derives
    @llvm
    def __new__(args: T, kwargs: K, M: Static[str], F: type, T: type, K: type) -> Partial[M, T, K, F]:
        %0 = insertvalue { {=T}, {=K} } undef, {=T} %args, 0
        %1 = insertvalue { {=T}, {=K} } %0, {=K} %kwargs, 1
        ret { {=T}, {=K} } %1

    # def __new__(M: Static[str], F: type) -> Partial[M, Tuple[Tuple], NamedTuple[0,Tuple], F]:
    #     return Partial.__new__((), NamedTuple[0,Tuple](), M=M, F=F)

    def __repr__(self):
        return __magic__.repr_partial(self)

    def __call__(self, *args, **kwargs):
        return self(*args, **kwargs)

    @property
    def __fn_name__(self):
        return F.T.__name__

    def __raw__(self):
        # TODO: better error message
        return F.T.__raw__()