# Copyright 2026 Marimo. All rights reserved.
"""Stub system for cache serialization."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from marimo._save.stubs.class_stub import ClassStub
from marimo._save.stubs.function_stub import FunctionStub
from marimo._save.stubs.lazy_stub import (
    LAZY_STUB_LOOKUP,
    BlobAsset,
    ReferenceStub,
)
from marimo._save.stubs.module_stub import ModuleStub
from marimo._save.stubs.pydantic_stub import PydanticStub
from marimo._save.stubs.stubs import (
    CUSTOM_STUBS,
    CustomStub,
    mro_lookup,
    register_stub,
)
from marimo._save.stubs.ui_element_stub import UIElementStub

if TYPE_CHECKING:
    from collections.abc import Callable

# Track which class names we've already attempted to register
_REGISTERED_NAMES: set[str] = set()

# Dictionary mapping fully qualified class names to registration functions
STUB_REGISTRATIONS: dict[str, Callable[[Any], None]] = {
    "pydantic.main.BaseModel": PydanticStub.register,
}


def maybe_register_stub(value: Any) -> bool:
    """Lazily register a stub for a value's type if not already registered.

    This allows us to avoid importing third-party packages until they're
    actually used in the cache. Walks the MRO to check if any parent class
    matches a registered stub type.

    Returns:
        True if the value's type is in CUSTOM_STUBS (either already registered
        or newly registered), False otherwise.
    """
    value_type = type(value)

    # Already registered in CUSTOM_STUBS
    if value_type in CUSTOM_STUBS:
        return True

    result = mro_lookup(value_type, STUB_REGISTRATIONS)
    if result is not None:
        cls_name, register_fn = result
        if cls_name not in _REGISTERED_NAMES:
            _REGISTERED_NAMES.add(cls_name)
            register_fn(value)
        return value_type in CUSTOM_STUBS

    return False


def maybe_get_custom_stub(value: Any) -> CustomStub | None:
    """Get the registered stub for a value's type, if any.

    Args:
        value: The value to get the stub for

    Returns:
        A stub instance if registered, None otherwise
    """
    # Fallback to custom cases
    if maybe_register_stub(value):
        value_type = type(value)
        if value_type in CUSTOM_STUBS:
            return CUSTOM_STUBS[value_type](value)
    return None


__all__ = [
    "CUSTOM_STUBS",
    "LAZY_STUB_LOOKUP",
    "BlobAsset",
    "ClassStub",
    "CustomStub",
    "FunctionStub",
    "ModuleStub",
    "ReferenceStub",
    "UIElementStub",
    "maybe_get_custom_stub",
    "maybe_register_stub",
    "mro_lookup",
    "register_stub",
]
