from types import FunctionType
from typing import Any, Never, assert_never

from sonolus.backend.ops import Op
from sonolus.script.array import Array
from sonolus.script.array_like import ArrayLike
from sonolus.script.debug import error, require
from sonolus.script.internal.context import ctx
from sonolus.script.internal.dict_impl import DictImpl
from sonolus.script.internal.impl import meta_fn, validate_value
from sonolus.script.internal.math_impls import MATH_BUILTIN_IMPLS, _trunc
from sonolus.script.internal.native import native_function
from sonolus.script.internal.random import RANDOM_BUILTIN_IMPLS
from sonolus.script.internal.range import Range
from sonolus.script.internal.tuple_impl import TupleImpl
from sonolus.script.internal.value import Value
from sonolus.script.iterator import (
    SonolusIterator,
    _EmptyIterator,
    _Enumerator,
    _FilteringIterator,
    _MappingIterator,
    _Zipper,
)
from sonolus.script.num import Num, _is_num
from sonolus.script.record import Record

_empty = object()


@meta_fn
def _isinstance(value, type_):
    value = validate_value(value)
    type_ = validate_value(type_)._as_py_()
    if type_ is dict:
        return isinstance(value, DictImpl)
    if type_ is tuple:
        return isinstance(value, TupleImpl)
    if type_ in {_int, _float, _bool}:
        raise TypeError("Instance check against int, float, or bool is not supported, use Num instead")
    if not (isinstance(type_, type) and (issubclass(type_, Value) or getattr(type_, "_allow_instance_check_", False))):
        raise TypeError(f"Unsupported type: {type_} for isinstance")
    return validate_value(isinstance(value, type_))


@meta_fn
def _len(value):
    from sonolus.backend.visitor import compile_and_call

    value = validate_value(value)
    if not hasattr(value, "__len__"):
        raise TypeError(f"object of type '{type(value).__name__}' has no len()")
    return compile_and_call(value.__len__)  # type: ignore


@meta_fn
def _enumerate(iterable, start=0):
    from sonolus.backend.visitor import compile_and_call

    iterable = validate_value(iterable)
    if isinstance(iterable, TupleImpl):
        return TupleImpl._accept_(tuple((start + i, value) for i, value in enumerate(iterable._as_py_(), start=start)))
    elif not hasattr(iterable, "__iter__"):
        raise TypeError(f"'{type(iterable).__name__}' object is not iterable")
    elif isinstance(iterable, ArrayLike):
        return compile_and_call(iterable._enumerate_, start)
    else:
        iterator = compile_and_call(iterable.__iter__)  # type: ignore
        if not isinstance(iterator, SonolusIterator):
            raise TypeError("Only subclasses of SonolusIterator are supported as iterators")
        return _Enumerator(0, start, iterator)


@meta_fn
def _reversed(iterable):
    from sonolus.backend.visitor import compile_and_call

    iterable = validate_value(iterable)
    if not isinstance(iterable, ArrayLike):
        raise TypeError(f"Unsupported type: {type(iterable)} for reversed")
    return compile_and_call(iterable.__reversed__)


@meta_fn
def _zip(*iterables, strict: bool = False):
    from sonolus.backend.visitor import compile_and_call
    from sonolus.script.containers import Pair

    if validate_value(strict)._as_py_():  # type: ignore
        raise NotImplementedError("Strict zipping is not supported")

    if not iterables:
        return _EmptyIterator()

    iterables = [validate_value(iterable) for iterable in iterables]
    if any(isinstance(iterable, TupleImpl) for iterable in iterables):
        if not all(isinstance(iterable, TupleImpl) for iterable in iterables):
            raise TypeError("Cannot mix tuples with other types in zip")
        return TupleImpl._accept_(tuple(zip(*(iterable.value for iterable in iterables), strict=False)))
    iterators = [compile_and_call(iterable.__iter__) for iterable in iterables]
    if not all(isinstance(iterator, SonolusIterator) for iterator in iterators):
        raise TypeError("Only subclasses of SonolusIterator are supported as iterators")
    v = iterators.pop()
    while iterators:
        v = Pair(iterators.pop(), v)
    return _Zipper(v)


@meta_fn
def _abs(value):
    from sonolus.backend.visitor import compile_and_call

    value = validate_value(value)
    if not hasattr(value, "__abs__"):
        raise TypeError(f"bad operand type for abs(): '{type(value).__name__}'")
    return compile_and_call(value.__abs__)  # type: ignore


def _identity(value):
    return value


@meta_fn
def _max(*args, default=_empty, key=None):
    from sonolus.backend.visitor import compile_and_call

    if key is None:
        key = _identity

    args = tuple(validate_value(arg) for arg in args)
    if len(args) == 0:
        raise ValueError("Expected at least one argument to max")
    elif len(args) == 1:
        (iterable,) = args
        if isinstance(iterable, ArrayLike):
            return compile_and_call(iterable._max_, key=key)
        elif isinstance(iterable, TupleImpl) and all(_is_num(v) for v in iterable.value):
            if len(iterable.value) == 0:
                if default is not _empty:
                    return default
                raise ValueError("max() arg is an empty sequence")
            return compile_and_call(Array(*iterable.value)._max_, key=key)
        elif isinstance(iterable, SonolusIterator):
            if not (default is _empty or Num._accepts_(default)):
                raise TypeError("default argument must be a number")
            return compile_and_call(
                _max_num_iterator,
                iterable,
                Num._accept_(default) if default is not _empty else None,
                key=key if key is not _identity else None,
            )
        else:
            raise TypeError(f"Unsupported type: {type(iterable)} for max")
    else:
        if default is not _empty:
            raise TypeError("default argument is not supported for max with multiple arguments")
        if not all(_is_num(arg) for arg in args):
            raise TypeError("Arguments to max must be numbers")
        if ctx():
            result = _max2(args[0], args[1], key=key)
            for arg in args[2:]:
                result = _max2(result, arg, key=key)
            return result
        else:
            return max(arg._as_py_() for arg in args)


def _max2(a, b, key=_identity):
    from sonolus.backend.visitor import compile_and_call

    a = validate_value(a)
    b = validate_value(b)
    if _is_num(a) and _is_num(b) and key == _identity:
        return compile_and_call(_max2_num, a, b)
    return compile_and_call(_max2_generic, a, b, key=key)


@native_function(Op.Max)
def _max2_num(a, b):
    if a > b:
        return a
    else:
        return b


def _max2_generic(a, b, key=_identity):
    if key(a) > key(b):
        return a
    else:
        return b


def _max_num_iterator(iterable, default, key):
    iterator = iterable.__iter__()  # noqa: PLC2801
    initial = iterator.next()
    if initial.is_nothing:
        require(default is not None, "default must be provided if the iterator is empty")
        return default
    if key is not None:
        result = initial.get_unsafe()
        best_key = key(result)
        for value in iterator:
            new_key = key(value)
            if new_key > best_key:
                result = value
                best_key = new_key
        return result
    else:
        result = initial.get_unsafe()
        for value in iterator:
            if value > result:  # noqa: PLR1730
                result = value
        return result


@meta_fn
def _min(*args, default=_empty, key=None):
    from sonolus.backend.visitor import compile_and_call

    if key is None:
        key = _identity

    args = tuple(validate_value(arg) for arg in args)
    if len(args) == 0:
        raise ValueError("Expected at least one argument to min")
    elif len(args) == 1:
        (iterable,) = args
        if isinstance(iterable, ArrayLike):
            return compile_and_call(iterable._min_, key=key)
        elif isinstance(iterable, TupleImpl) and all(_is_num(v) for v in iterable.value):
            if len(iterable.value) == 0:
                if default is not _empty:
                    return default
                raise ValueError("min() arg is an empty sequence")
            return compile_and_call(Array(*iterable.value)._min_, key=key)
        elif isinstance(iterable, SonolusIterator):
            if not (default is _empty or Num._accepts_(default)):
                raise TypeError("default argument must be a number")
            return compile_and_call(
                _min_num_iterator,
                iterable,
                Num._accept_(default) if default is not _empty else None,
                key=key if key is not _identity else None,
            )
        else:
            raise TypeError(f"Unsupported type: {type(iterable)} for min")
    else:
        if default is not _empty:
            raise TypeError("default argument is not supported for min with multiple arguments")
        if not all(_is_num(arg) for arg in args):
            raise TypeError("Arguments to min must be numbers")
        if ctx():
            result = _min2(args[0], args[1], key=key)
            for arg in args[2:]:
                result = _min2(result, arg, key=key)
            return result
        else:
            return min(arg._as_py_() for arg in args)


def _min2(a, b, key=_identity):
    from sonolus.backend.visitor import compile_and_call

    a = validate_value(a)
    b = validate_value(b)
    if _is_num(a) and _is_num(b) and key == _identity:
        return compile_and_call(_min2_num, a, b)
    return compile_and_call(_min2_generic, a, b, key=key)


@native_function(Op.Min)
def _min2_num(a, b):
    if a < b:
        return a
    else:
        return b


def _min2_generic(a, b, key=_identity):
    if key(a) < key(b):
        return a
    else:
        return b


def _min_num_iterator(iterable, default, key):
    iterator = iterable.__iter__()  # noqa: PLC2801
    initial = iterator.next()
    if initial.is_nothing:
        require(default is not None, "default must be provided if the iterator is empty")
        return default
    if key is not None:
        result = initial.get_unsafe()
        best_key = key(result)
        for value in iterator:
            new_key = key(value)
            if new_key < best_key:
                result = value
                best_key = new_key
        return result
    else:
        result = initial.get_unsafe()
        for value in iterator:
            if value < result:  # noqa: PLR1730
                result = value
        return result


@meta_fn
def _callable(value):
    return callable(value)


def _map(fn, iterable, *iterables):
    if len(iterables) == 0:
        return _MappingIterator(fn, iterable.__iter__())  # noqa: PLC2801
    return _MappingIterator(lambda args: fn(*args), zip(iterable, *iterables))  # noqa: B905


def _filter(fn, iterable):
    if fn is None:
        fn = _identity
    return _FilteringIterator(fn, iterable.__iter__())  # noqa: PLC2801


@meta_fn
def _int(value=0):
    value = validate_value(value)
    if not _is_num(value):
        raise TypeError("Only numeric arguments to int() are supported")
    return _trunc(value)


@meta_fn
def _float(value=0.0):
    value = validate_value(value)
    if not _is_num(value):
        raise TypeError("Only numeric arguments to float() are supported")
    return value


def _bool(value=False):
    # Relies on the compiler to perform the conversion in a boolean context
    if value:  # noqa: SIM103
        return True
    else:
        return False


_int._type_mapping_ = Num  # type: ignore
_float._type_mapping_ = Num  # type: ignore
_bool._type_mapping_ = Num  # type: ignore


def _any(iterable):
    for value in iterable:  # noqa: SIM110
        if value:
            return True
    return False


def _all(iterable):
    for value in iterable:  # noqa: SIM110
        if not value:
            return False
    return True


def _sum(iterable, /, start=0):
    for value in iterable:
        start += value
    return start


def _next(iterator):
    require(isinstance(iterator, SonolusIterator), "Only subclasses of SonolusIterator are supported as iterators")
    value = iterator.next()
    if value.is_some:
        return value.get_unsafe()
    error("Iterator has been exhausted")


def _iter(iterable):
    return iterable.__iter__()  # type: ignore # noqa: PLC2801


@meta_fn
def _super(*args):
    """Get the super class of a class or instance."""
    return super(*(arg._as_py_() if arg._is_py_() else arg for arg in args))


@meta_fn
def _hasattr(obj: Any, name: str) -> bool:
    from sonolus.script.internal.constant import ConstantValue
    from sonolus.script.internal.descriptor import SonolusDescriptor

    name = validate_value(name)._as_py_()
    if isinstance(obj, ConstantValue):
        # Unwrap so we can access fields
        obj = obj._as_py_()
    descriptor = None
    for cls in type.mro(type(obj)):
        descriptor = cls.__dict__.get(name, None)
        if descriptor is not None:
            break
    # We want to mirror what getattr supports and fail fast if a future getattr would fail.
    match descriptor:
        case None:
            return hasattr(obj, name)
        case property() | SonolusDescriptor() | FunctionType() | classmethod() | staticmethod():
            return True
        case non_descriptor if not hasattr(non_descriptor, "__get__"):
            return True
        case _:
            raise TypeError(f"Unsupported field or descriptor {name}")


@meta_fn
def _getattr(obj: Any, name: str, default=_empty) -> Any:
    from sonolus.backend.visitor import compile_and_call
    from sonolus.script.internal.constant import ConstantValue
    from sonolus.script.internal.descriptor import SonolusDescriptor

    name = validate_value(name)._as_py_()
    if isinstance(obj, ConstantValue):
        # Unwrap so we can access fields
        obj = obj._as_py_()
    descriptor = None
    for cls in type.mro(type(obj)):
        descriptor = cls.__dict__.get(name, None)
        if descriptor is not None:
            break
    match descriptor:
        case property(fget=getter):
            return compile_and_call(getter, obj)
        case SonolusDescriptor() | FunctionType() | classmethod() | staticmethod() | None:
            return validate_value(getattr(obj, name) if default is _empty else getattr(obj, name, default))
        case non_descriptor if not hasattr(non_descriptor, "__get__"):
            return validate_value(getattr(obj, name) if default is _empty else getattr(obj, name, default))
        case _:
            raise TypeError(f"Unsupported field or descriptor {name}")


@meta_fn
def _setattr(obj: Any, name: str, value: Any):
    from sonolus.backend.visitor import compile_and_call
    from sonolus.script.internal.descriptor import SonolusDescriptor

    name = validate_value(name)._as_py_()
    if obj._is_py_():
        obj = obj._as_py_()
    descriptor = getattr(type(obj), name, None)
    match descriptor:
        case property(fset=setter):
            if setter is None:
                raise AttributeError(f"Cannot set attribute {name} because property has no setter")
            compile_and_call(setter, obj, value)
        case SonolusDescriptor():
            setattr(obj, name, value)
        case _:
            raise TypeError(f"Unsupported field or descriptor {name}")


class _Type(Record):
    @meta_fn
    def __call__(self, value, /):
        value = validate_value(value)
        if value._is_py_():
            value = value._as_py_()
        return validate_value(type(value))

    def __getitem__(self, item):
        return self


_type = _Type()


@meta_fn
def _assert_never(arg: Never, /):
    error("Expected code to be unreachable")


# classmethod, property, staticmethod are supported as decorators, but not within functions

BUILTIN_IMPLS = {
    id(abs): _abs,
    id(all): _all,
    id(any): _any,
    id(bool): _bool,
    id(callable): _callable,
    id(enumerate): _enumerate,
    id(filter): _filter,
    id(float): _float,
    id(getattr): _getattr,
    id(hasattr): _hasattr,
    id(int): _int,
    id(isinstance): _isinstance,
    id(iter): _iter,
    id(len): _len,
    id(map): _map,
    id(max): _max,
    id(min): _min,
    id(next): _next,
    id(range): Range,
    id(reversed): _reversed,
    id(setattr): _setattr,
    id(sum): _sum,
    id(super): _super,
    id(type): _type,
    id(zip): _zip,
    id(assert_never): _assert_never,
    **MATH_BUILTIN_IMPLS,  # Includes round
    **RANDOM_BUILTIN_IMPLS,
}
