"""
This module provides thread-safe code context management for TorchDynamo using weak references.

The CodeContextDict class maintains a mapping between Python code objects and their associated
context data, using weak references to automatically clean up entries when code objects are
garbage collected. This prevents memory leaks while allowing context data to be associated
with code objects throughout their lifecycle.

Key features:
- Thread-safe context storage and retrieval
- Automatic cleanup using weak references
- Safe context management for Python code objects
- Memory-leak prevention

Example usage:
    code_obj = compile('x = 1', '<string>', 'exec')

    # Store context
    context = code_context.get_context(code_obj)
    context['metadata'] = {'optimized': True}

    # Retrieve context
    if code_context.has_context(code_obj):
        ctx = code_context.get_context(code_obj)
        # Use context data...

    # Remove context
    ctx = code_context.pop_context(code_obj)
"""

import types
from typing import Any

from .utils import ExactWeakKeyDictionary


class CodeContextDict:
    def __init__(self) -> None:
        self.code_context: ExactWeakKeyDictionary = ExactWeakKeyDictionary()

    def has_context(self, code: types.CodeType) -> bool:
        return code in self.code_context

    def get_context(self, code: types.CodeType) -> dict[str, Any]:
        ctx = self.code_context.get(code)
        if ctx is None:
            ctx = {}
            self.code_context[code] = ctx
        return ctx

    def pop_context(self, code: types.CodeType) -> dict[str, Any]:
        ctx = self.get_context(code)
        self.code_context._remove_id(id(code))
        return ctx

    def clear(self) -> None:
        self.code_context.clear()


code_context: CodeContextDict = CodeContextDict()