import json
from enum import Enum
from typing import Any, Dict, List, Union, Type
from pathlib import Path

try:
    import pygame

    PYGAME_AVAILABLE = True
except ImportError:
    PYGAME_AVAILABLE = False


class EnumSerializer:
    """
    A class to serialize and deserialize Enum classes and collections to/from JSON files.
    Handles dynamic Enum structures, various collection types, and Pygame keys.
    """

    # Cache for Pygame key mappings
    _pygame_key_to_name = None
    _pygame_name_to_key = None

    @classmethod
    def _init_pygame_mappings(cls):
        """Initialize Pygame key mappings if not already done."""
        if not PYGAME_AVAILABLE or cls._pygame_key_to_name is not None:
            return

        cls._pygame_key_to_name = {}
        cls._pygame_name_to_key = {}

        # Get all K_* constants from pygame
        for attr_name in dir(pygame):
            if attr_name.startswith('K_'):
                key_value = getattr(pygame, attr_name)
                cls._pygame_key_to_name[key_value] = attr_name
                cls._pygame_name_to_key[attr_name] = key_value

    @classmethod
    def _is_pygame_key(cls, value: Any) -> bool:
        """Check if a value is a Pygame key constant."""
        if not PYGAME_AVAILABLE or not isinstance(value, int):
            return False

        cls._init_pygame_mappings()
        return value in cls._pygame_key_to_name

    @classmethod
    def _serialize_pygame_key(cls, key_value: int) -> Dict[str, Any]:
        """Serialize a Pygame key to human-readable format."""
        cls._init_pygame_mappings()
        return {
            '__type__': 'PygameKey',
            'name': cls._pygame_key_to_name[key_value],
            'value': key_value  # Keep numeric value as backup
        }

    @staticmethod
    def _serialize_enum_class(enum_class: Type[Enum]) -> Dict[str, Any]:
        """Serialize an Enum class definition to a dictionary."""
        return {
            '__type__': 'EnumClass',
            'name': enum_class.__name__,
            'members': {member.name: member.value for member in enum_class}
        }

    @staticmethod
    def _serialize_enum_instance(enum_instance: Enum) -> Dict[str, Any]:
        """Serialize an Enum instance to a dictionary."""
        return {
            '__type__': 'EnumInstance',
            'class_name': enum_instance.__class__.__name__,
            'name': enum_instance.name,
            'value': enum_instance.value
        }

    @staticmethod
    def _serialize_value(value: Any) -> Any:
        """Recursively serialize a value, handling Enums, Pygame keys, and collections."""
        # Check for Pygame key first
        if EnumSerializer._is_pygame_key(value):
            return EnumSerializer._serialize_pygame_key(value)
        elif isinstance(value, Enum):
            return EnumSerializer._serialize_enum_instance(value)
        elif isinstance(value, dict):
            return {k: EnumSerializer._serialize_value(v) for k, v in value.items()}
        elif isinstance(value, (list, tuple, set)):
            serialized = [EnumSerializer._serialize_value(item) for item in value]
            return {
                '__type__': type(value).__name__,
                'items': serialized
            }
        else:
            return value

    @staticmethod
    def save_to_file(data: Dict[str, Any], filename: str) -> None:
        """
        Save data (Enums, Pygame keys, and collections) to a JSON file.

        Args:
            data: Dictionary containing the data to save. Keys are variable names,
                  values can be Enum classes, Enum instances, Pygame keys, or collections.
            filename: Path to the JSON file where data will be saved.

        Example:
            serializer = EnumSerializer()
            data = {
                'color_enum': Color,
                'current_color': Color.RED,
                'color_list': [Color.RED, Color.BLUE],
                'settings': {'theme': Color.DARK, 'count': 5},
                'move_left': pygame.K_LEFT,
                'controls': {'left': pygame.K_LEFT, 'right': pygame.K_RIGHT}
            }
            serializer.save_to_file(data, 'config.json')
        """
        serialized_data = {}

        for key, value in data.items():
            if isinstance(value, type) and issubclass(value, Enum):
                # Serialize Enum class definition
                serialized_data[key] = EnumSerializer._serialize_enum_class(value)
            else:
                # Serialize value (could be Enum instance, Pygame key, or collection)
                serialized_data[key] = EnumSerializer._serialize_value(value)

        Path(filename).parent.mkdir(parents=True, exist_ok=True)

        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(serialized_data, f, indent=2, ensure_ascii=False)

    @staticmethod
    def _deserialize_enum_class(data: Dict[str, Any]) -> Type[Enum]:
        """Recreate an Enum class from serialized data."""
        name = data['name']
        members = data['members']
        return Enum(name, members)

    @classmethod
    def _deserialize_value(cls, value: Any, enum_registry: Dict[str, Type[Enum]]) -> Any:
        """Recursively deserialize a value, handling Enums, Pygame keys, and collections."""
        if isinstance(value, dict):
            if '__type__' in value:
                type_name = value['__type__']

                if type_name == 'PygameKey':
                    # Restore Pygame key
                    if PYGAME_AVAILABLE:
                        cls._init_pygame_mappings()
                        key_name = value['name']
                        if key_name in cls._pygame_name_to_key:
                            return cls._pygame_name_to_key[key_name]
                    # Fallback to numeric value if Pygame not available
                    return value['value']

                elif type_name == 'EnumInstance':
                    # Restore Enum instance
                    class_name = value['class_name']
                    if class_name in enum_registry:
                        enum_class = enum_registry[class_name]
                        return enum_class[value['name']]
                    else:
                        # Return raw data if Enum class not found
                        return value

                elif type_name in ('list', 'tuple', 'set'):
                    # Restore collection
                    items = [cls._deserialize_value(item, enum_registry)
                             for item in value['items']]
                    if type_name == 'tuple':
                        return tuple(items)
                    elif type_name == 'set':
                        return set(items)
                    else:
                        return items
            else:
                # Regular dictionary
                return {k: cls._deserialize_value(v, enum_registry)
                        for k, v in value.items()}
        else:
            return value

    @staticmethod
    def load_from_file(filename: str) -> Dict[str, Any]:
        """
        Load data (Enums, Pygame keys, and collections) from a JSON file.

        Args:
            filename: Path to the JSON file to load.

        Returns:
            Dictionary with deserialized data. Enum classes, instances, and Pygame keys are restored.

        Example:
            serializer = EnumSerializer()
            data = serializer.load_from_file('config.json')
            Color = data['color_enum']
            current_color = data['current_color']
            left_key = data['move_left']  # Will be pygame.K_LEFT
        """
        with open(filename, 'r', encoding='utf-8') as f:
            serialized_data = json.load(f)

        # First pass: restore all Enum class definitions
        enum_registry = {}
        for key, value in serialized_data.items():
            if isinstance(value, dict) and value.get('__type__') == 'EnumClass':
                enum_class = EnumSerializer._deserialize_enum_class(value)
                enum_registry[enum_class.__name__] = enum_class

        # Second pass: deserialize all values
        deserialized_data = {}
        for key, value in serialized_data.items():
            if isinstance(value, dict) and value.get('__type__') == 'EnumClass':
                # Return the Enum class itself
                deserialized_data[key] = enum_registry[value['name']]
            else:
                # Deserialize value
                deserialized_data[key] = EnumSerializer._deserialize_value(value, enum_registry)

        return deserialized_data


if __name__ == '__main__':
    pass
