from __future__ import annotations
import copy
import sys
from typing import Any, Dict, Optional, Union, TYPE_CHECKING

if TYPE_CHECKING:
    from cosapp.base import System


def get_state(system: System) -> dict:
    """Export system state (inputs and outputs) into a dictionary
    for quick serialization, without type check.
    
    Returns:
    --------
    Recursive dictionary of the kind:
    {
        'ports: {
            port_name: {
                varname: value,
                ...
            },
            ...
        },
        'children': {
            child_name: {
                'ports': {...},
                'children: {...},
            },
            ...
        },
    }
    """
    from cosapp.base import System

    def recursive_search(system: System, state: dict) -> None:
        state['ports'] = port_data = {}
        for port in system.ports():
            data = {
                var: copy.deepcopy(value)
                for var, value in port.items()
            }
            if data:
                port_data[port.name] = data
        if system.children:
            state['children'] = child_data = {}
            for child in system.children.values():
                child_data[child.name] = child_state = {}
                recursive_search(child, child_state)

    state = {}
    recursive_search(system, state)

    return state


def set_state(system: System, state: dict) -> None:
    """Import system state (inputs and outputs) from
    a dictionary generated by `get_state`.
    """
    for name, values in state['ports'].items():
        port = getattr(system, name)
        for name, value in values.items():
            setattr(port, name, copy.deepcopy(value))
        # system[name].set_values(**copy.deepcopy(values))

    child_data = state.get('children', {})

    for name, child_state in child_data.items():
        set_state(system[name], child_state)


if sys.version_info[1] <11:
    def object__getstate__(obj: Any) -> Union[Dict[str, Any], tuple[Optional[Dict[str, Any]], Dict[str, Any]]]:
        """Creates a state of an object.

        The state may take various forms depending on the object, see
        https://docs.python.org/3/library/pickle.html#object.__getstate__
        for further details.
        
        Parameters
        ----------
        obj: Any
            Object from which the state must be constructed

        Returns
        -------
        Union[Dict[str, Any], tuple[Optional[Dict[str, Any]], Dict[str, Any]]]:
            state
        """
        def is_slot(slot, cls):
            not_magic_slot = not (slot.startswith("__") and slot.endswith("__"))
            cls_slot = hasattr(cls, slot)
            priv_cls_slot = hasattr(cls, f"_{cls.__name__}{slot}")
            return not_magic_slot and (cls_slot or priv_cls_slot)

        def prefix_private_slot(slot, cls):
            if slot.startswith("__") and not slot.endswith("__"):
                return f"_{cls.__name__}{slot}"
            return slot

        if hasattr(obj, "__slots__"):
            slots_names = [prefix_private_slot(slot, cls) for cls in obj.__class__.__mro__ for slot in getattr(cls, '__slots__', []) if is_slot(slot, cls)]
            slots = {slot: getattr(obj, slot) for slot in slots_names}

            if slots:
                return getattr(obj, "__dict__", None), slots

        if hasattr(obj, "__dict__"):
            return obj.__dict__
else:
    object__getstate__ = object.__getstate__