from __future__ import annotations

import inspect
from abc import ABCMeta
from collections.abc import Callable, Iterable
from functools import wraps
from inspect import getmro
from typing import Any, ClassVar, Protocol, Self, TypeVar, dataclass_transform, get_origin

from sonolus.backend.place import BlockPlace
from sonolus.script.internal.context import ctx
from sonolus.script.internal.descriptor import SonolusDescriptor
from sonolus.script.internal.generic import (
    GenericValue,
    accept_and_infer_types,
    validate_and_resolve_type,
    validate_concrete_type,
    validate_type_spec_with_extras,
)
from sonolus.script.internal.meta_fn import meta_fn
from sonolus.script.internal.value import BackingSource, DataValue, Value
from sonolus.script.num import Num


# Protocol metaclass is a subclass of ABCMeta, but let's make ABCMeta explicit for clarity
class RecordMeta(type(Protocol), ABCMeta):
    @meta_fn
    def __pos__[T](cls: type[T]) -> T:
        """Create a zero-initialized record instance."""
        return cls._zero_()


@dataclass_transform(eq_default=True)
class Record(GenericValue, metaclass=RecordMeta):
    """Base class for user-defined data structures.

    Usage:
        A regular record:
        ```python
        class MyRecord(Record):
            field1: int
            field2: bool
        ```

        A generic record:
        ```python
        class MyGenericRecord[T, U](Record):
            field1: T
            field2: U
        ```

        Creating an instance:
        ```python
        record = MyRecord(field1=42, field2=True)
        record_2 = MyGenericRecord[int, int](field1=42, field2=100)
        record_3 = MyGenericRecord(field1=42, field2=100)  # Type arguments can be inferred
        record_4 = +MyRecord  # Create a zero-initialized record
        record_5 = +MyGenericRecord[int, int]
        ```

        Copying a record:
        ```python
        record_copy = +record
        ```
    """

    _value_: dict[str, Value]
    _fields_: ClassVar[list[_RecordField] | None] = None
    _constructor_signature_: ClassVar[inspect.Signature]
    _is_frozen_: bool = False

    @classmethod
    @meta_fn
    def frozen[**P](cls: Callable[P, Self], /, *args: P.args, **kwargs: P.kwargs) -> Self:
        """Create a frozen record instance.

        Numeric fields in the record will not be writable. However, nested values such as other records may remain
        mutable.

        Args:
            *args: Positional arguments for the record constructor.
            **kwargs: Keyword arguments for the record constructor.

        Returns:
            A frozen record instance.
        """
        # We override __new__ to allow changing to the parameterized version
        if cls._constructor_signature_ is None:
            raise TypeError(f"Cannot instantiate {cls.__name__}")
        bound = cls._constructor_signature_.bind(*args, **kwargs)
        bound.apply_defaults()
        values = {}
        type_vars = {}
        for field in cls._fields_:
            value = bound.arguments[field.name]
            value = accept_and_infer_types(field.type, value, type_vars)
            values[field.name] = value._get_readonly_()
        for type_param in cls.__type_params__:
            if type_param not in type_vars:
                raise TypeError(f"Type parameter {type_param} cannot be inferred and must be provided explicitly")
        type_args = tuple(type_vars[type_param] for type_param in cls.__type_params__)
        if cls._type_args_ is not None:
            parameterized = cls
        else:
            parameterized = cls[type_args]
        result: cls = object.__new__(parameterized)  # type: ignore
        result._value_ = values
        result._is_frozen_ = True
        return result

    @classmethod
    def _validate_type_args_(cls, args: tuple[Any, ...]) -> tuple[Any, ...]:
        if cls._fields_ is None:
            raise TypeError("Base Record class cannot have type arguments")
        return super()._validate_type_args_(args)

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__()
        is_parameterizing = cls._type_args_ is not None and all(
            getattr(parent, "_type_args_", None) is None for parent in getmro(cls)[1:]
        )
        if is_parameterizing:
            fields = []
            offset = 0
            for generic_field in cls._fields_:
                resolved_type = validate_and_resolve_type(generic_field.type, cls._type_vars_to_args_)
                resolved_type = validate_concrete_type(resolved_type)
                field = _RecordField(
                    generic_field.name, resolved_type, generic_field.index, offset, generic_field.final
                )
                fields.append(field)
                setattr(cls, field.name, field)
                offset += resolved_type._size_()
            cls._fields_ = fields
            return
        is_inheriting_from_existing_record_class = cls._fields_ is not None
        if is_inheriting_from_existing_record_class and not is_parameterizing:
            # The main reason this is disallowed is that subclasses wouldn't be substitutable for their parent classes
            # Assignment of a subclass instance to a variable of the parent class would either be disallowed or would
            # require object slicing. Either way, it could lead to confusion.
            # Dealing with generic supertypes is also tricky, so it isn't really worth the effort to support this.
            raise TypeError("Subclassing of a Record is not supported")

        hints = inspect.get_annotations(cls, eval_str=True)
        fields = []
        params = []
        index = 0
        offset = 0
        for name, hint in hints.items():
            try:
                if name not in cls.__annotations__:
                    continue
                if hint is ClassVar or get_origin(hint) is ClassVar:
                    continue
                if hasattr(cls, name):
                    raise TypeError("Default values are not supported for Record fields")
                type_info = validate_type_spec_with_extras(hint)
                type_ = type_info.value
                fields.append(_RecordField(name, type_, index, offset, type_info.final))
                if isinstance(type_, type) and issubclass(type_, Value) and type_._is_concrete_():
                    offset += type_._size_()
                setattr(cls, name, fields[-1])
                index += 1
                params.append(
                    inspect.Parameter(
                        name,
                        inspect.Parameter.POSITIONAL_OR_KEYWORD,
                        annotation=type_,
                    )
                )
            except Exception as e:
                raise TypeError(f"Error processing field '{name}' of Record {cls.__name__}: {e}") from e

        cls._parameterized_ = {}
        cls._fields_ = fields
        cls._constructor_signature_ = inspect.Signature(params)

        _add_inplace_ops(cls)

        cls.__match_args__ = tuple(field.name for field in fields)

        if len(getattr(cls, "__type_params__", ())) == 0:
            # Make the class behave as the parameterized version
            cls._type_args_ = ()
            cls._type_vars_to_args_ = {}
            cls._parameterized_[()] = cls

    def __new__(cls, *args, **kwargs):
        # We override __new__ to allow changing to the parameterized version
        if cls._constructor_signature_ is None:
            raise TypeError(f"Cannot instantiate {cls.__name__}")
        bound = cls._constructor_signature_.bind(*args, **kwargs)
        bound.apply_defaults()
        values = {}
        type_vars = {}
        for field in cls._fields_:
            value = bound.arguments[field.name]
            value = accept_and_infer_types(field.type, value, type_vars)
            if field.final:
                values[field.name] = value._get_readonly_()
            else:
                values[field.name] = value._get_()
        for type_param in cls.__type_params__:
            if type_param not in type_vars:
                raise TypeError(f"Type parameter {type_param} cannot be inferred and must be provided explicitly")
        type_args = tuple(type_vars[type_param] for type_param in cls.__type_params__)
        if cls._type_args_ is not None:
            parameterized = cls
        else:
            parameterized = cls[type_args]
        result: cls = object.__new__(parameterized)  # type: ignore
        result._value_ = values
        return result

    def __init__(self, *args, **kwargs):
        # Initialization is done in __new__ and other methods
        pass

    @classmethod
    def _raw(cls, **kwargs) -> Self:
        result = object.__new__(cls)
        result._value_ = kwargs
        return result

    @classmethod
    def _unchecked(cls, **kwargs) -> Self:
        # Skips most validation, generally for internal use in frequently-called methods for performance reasons
        result = object.__new__(cls)
        for k, v in kwargs.items():
            if isinstance(v, int | float):
                kwargs[k] = Num._accept_(v)
        result._value_ = kwargs
        return result

    @classmethod
    def _size_(cls) -> int:
        return sum(field.type._size_() for field in cls._fields_)

    @classmethod
    def _is_value_type_(cls) -> bool:
        return False

    @classmethod
    def _from_backing_source_(cls, source: BackingSource) -> Self:
        result = object.__new__(cls)
        result._value_ = {
            field.name: field.type._from_backing_source_(
                lambda offset, field_offset=field.offset: source((Num(offset) + Num(field_offset)).ir())  # type: ignore
            )
            for field in cls._fields_
        }
        return result

    @classmethod
    def _from_place_(cls, place: BlockPlace) -> Self:
        result = object.__new__(cls)
        result._value_ = {field.name: field.type._from_place_(place.add_offset(field.offset)) for field in cls._fields_}
        return result

    @classmethod
    def _accepts_(cls, value: Any) -> bool:
        return issubclass(type(value), cls)

    @classmethod
    def _accept_(cls, value: Any) -> Self:
        if not cls._accepts_(value):
            raise TypeError(f"Cannot accept value {value} as {cls.__name__}")
        return value

    def _is_py_(self) -> bool:
        return all(value._is_py_() for value in self._value_.values())

    def _as_py_(self) -> Self:
        if not self._is_py_():
            raise ValueError("Not a python value")
        return self

    @classmethod
    def _from_list_(cls, values: Iterable[DataValue]) -> Self:
        iterator = iter(values)
        return cls._raw(**{field.name: field.type._from_list_(iterator) for field in cls._fields_})

    def _to_list_(self, level_refs: dict[Any, str] | None = None) -> list[DataValue | str]:
        result = []
        for field in self._fields_:
            result.extend(self._value_[field.name]._to_list_(level_refs))
        return result

    @classmethod
    def _flat_keys_(cls, prefix: str) -> list[str]:
        result = []
        for field in cls._fields_:
            result.extend(field.type._flat_keys_(f"{prefix}.{field.name}"))
        if len(result) == 1:
            return [prefix]
        return result

    def _get_(self) -> Self:
        return self

    def _set_(self, value: Any):
        raise TypeError("Record does not support _set_")

    def _copy_from_(self, value: Any):
        value = self._accept_(value)
        for field in self._fields_:
            field.__set__(self, field.__get__(value))

    def _copy_(self) -> Self:
        return type(self)._raw(**{field.name: self._value_[field.name]._copy_() for field in self._fields_})

    @classmethod
    def _alloc_(cls) -> Self:
        # Compared to using the constructor, this avoids unnecessary _get_ calls
        result = object.__new__(cls)
        result._value_ = {field.name: field.type._alloc_() for field in cls._fields_}
        return result

    @classmethod
    def _zero_(cls) -> Self:
        result = object.__new__(cls)
        result._value_ = {field.name: field.type._zero_() for field in cls._fields_}
        return result

    def __str__(self):
        return f"{self.__class__.__name__}({
            ', '.join(f'{field.name}={field.get_internal(self)}' for field in self._fields_)
        })"

    def __repr__(self):
        return f"{self.__class__.__name__}({
            ', '.join(f'{field.name}={field.get_internal(self)!r}' for field in self._fields_)
        })"

    @meta_fn
    def __eq__(self, other: Any) -> bool:
        if not isinstance(other, type(self)):
            return False
        result: Num = Num._accept_(True)
        for field in self._fields_:
            result = result.and_(field.__get__(self) == field.__get__(other))
        return result

    @meta_fn
    def __ne__(self, other: Any) -> bool:
        if not isinstance(other, type(self)):
            return True
        result: Num = Num._accept_(False)
        for field in self._fields_:
            result = result.or_(field.__get__(self) != field.__get__(other))
        return result

    def __hash__(self):
        return hash(tuple(field.__get__(self) for field in self._fields_))

    @meta_fn
    def __pos__(self) -> Self:
        """Return a copy of the record."""
        return self._copy_()

    @classmethod
    @meta_fn
    def type_var_value(cls, var: TypeVar, /) -> Any:
        """Return the value of a type variable.

        Args:
            var: The type variable to get the value of.

        Returns:
            The value of the type variable.
        """
        return super().type_var_value(var)

    @classmethod
    def _validate_parameterized_(cls):
        pass

    @classmethod
    def _get_parameterized(cls, args: tuple[Any, ...]) -> type[Self]:
        result = super()._get_parameterized(args)
        result._validate_parameterized_()
        return result


class _RecordField(SonolusDescriptor):
    def __init__(self, name: str, type_: type[Value] | Any, index: int, offset: int, final: bool):
        self.name = name
        self.type = type_
        self.index = index
        self.offset = offset
        self.final = final

    def get_internal(self, instance: Record) -> Value:
        return instance._value_[self.name]

    def __get__(self, instance: Record | None, owner=None):
        if instance is None:
            return self
        if instance._is_frozen_ or self.final:
            result = instance._value_[self.name]
        else:
            result = instance._value_[self.name]._get_readonly_()
        if ctx():
            return result
        else:
            return result._as_py_()

    def __set__(self, instance: Record, value):
        if instance._is_frozen_:
            raise TypeError("Cannot set fields of a frozen Record")
        if self.final:
            raise TypeError("Cannot set a final field")
        value = self.type._accept_(value)
        if self.type._is_value_type_():
            instance._value_[self.name]._set_(value)
        else:
            instance._value_[self.name]._copy_from_(value)


_ops_to_inplace_ops = {
    "__add__": "__iadd__",
    "__sub__": "__isub__",
    "__mul__": "__imul__",
    "__truediv__": "__itruediv__",
    "__floordiv__": "__ifloordiv__",
    "__mod__": "__imod__",
    "__pow__": "__ipow__",
    "__lshift__": "__ilshift__",
    "__rshift__": "__irshift__",
    "__or__": "__ior__",
    "__xor__": "__ixor__",
    "__and__": "__iand__",
    "__matmul__": "__imatmul__",
}


def _add_inplace_ops(cls):
    for op, inplace_op in _ops_to_inplace_ops.items():
        # Note that hasattr doesn't work, since it'll return True for metaclass methods
        if op in cls.__dict__ and not hasattr(cls, inplace_op):
            setattr(cls, inplace_op, _make_inplace_op(op, getattr(cls, op)))
    return cls


def _make_inplace_op(op: str, orig_fn):
    @meta_fn
    @wraps(orig_fn)
    def inplace_op(self, other):
        _compiler_internal_ = True  # noqa: F841
        self._copy_from_(getattr(self, op)(other))
        return self

    return inplace_op
