# Modified from: keras/src/backend/common/global_state.py
# Original authors: François Chollet et al. (Keras Team)
# License Apache 2.0: (c) 2025 Yoan Sallami (Synalinks Team)

import gc
import threading

from synalinks.src.api_export import synalinks_export

GLOBAL_STATE_TRACKER = threading.local()
GLOBAL_SETTINGS_TRACKER = threading.local()


def set_global_attribute(name, value):
    setattr(GLOBAL_STATE_TRACKER, name, value)


def get_global_attribute(name, default=None, set_to_default=False):
    attr = getattr(GLOBAL_STATE_TRACKER, name, None)
    if attr is None and default is not None:
        attr = default
        if set_to_default:
            set_global_attribute(name, attr)
    return attr


@synalinks_export(
    [
        "synalinks.utils.clear_session",
        "synalinks.backend.clear_session",
        "synalinks.clear_session",
    ]
)
def clear_session(free_memory=True):
    """Resets all state generated by synalinks.

    synalinks manages a global state, which it uses to implement the Functional
    model-building API and to uniquify autogenerated modules names.

    If you are creating many models in a loop, this global state will consume
    an increasing amount of memory over time, and you may want to clear it.
    Calling `clear_session()` releases the global state: this helps avoid
    clutter from old programs and modules, especially when memory is limited.

    Args:
        free_memory (bool): Whether to call Python garbage collection.
            It's usually a good practice to call it to make sure
            memory used by deleted objects is immediately freed.
            However, it may take a few seconds to execute, so
            when using `clear_session()` in a short loop,
            you may want to skip it.
    """
    global GLOBAL_STATE_TRACKER
    global GLOBAL_SETTINGS_TRACKER

    GLOBAL_STATE_TRACKER = threading.local()
    GLOBAL_SETTINGS_TRACKER = threading.local()

    if free_memory:
        # Manually trigger garbage collection.
        gc.collect()
