from lazyregistry import Registry

registry = Registry(name="plugins")

# Register by import string (lazy - imported on access)
registry["json"] = "json:dumps"

# Register by instance (immediate - already imported)
import pickle
registry["pickle"] = pickle.dumps

# Import happens here
serializer = registry["json"]

from lazyregistry import Registry

PLUGINS = Registry(name="plugins")

def plugin(name: str):
    def decorator(cls):
        PLUGINS[name] = cls
        return cls
    return decorator

@plugin("uppercase")
class UppercasePlugin:
    def process(self, text: str) -> str:
        return text.upper()

# Execute plugins
PluginManager.execute("uppercase", "hello")  # "HELLO"
PluginManager.pipeline("hello", "uppercase", "reverse")  # "OLLEH"

from lazyregistry import NAMESPACE
from lazyregistry.pretrained import AutoRegistry, PretrainedConfig, PretrainedMixin

# Each model has its own config with type identifier
class BertConfig(PretrainedConfig):
    model_type: str = "bert"
    hidden_size: int = 768

class GPT2Config(PretrainedConfig):
    model_type: str = "gpt2"
    hidden_size: int = 768

# Base model class
class BaseModel(PretrainedMixin):
    config_class = PretrainedConfig

class AutoModel(AutoRegistry):
    registry = NAMESPACE["models"]
    config_class = PretrainedConfig
    type_key = "model_type"

# Register with decorator - models inherit from BaseModel
@AutoModel.register_module("bert")
class BertModel(BaseModel):
    config_class = BertConfig

# Or register directly
AutoModel.registry["gpt2"] = "transformers:GPT2Model"  # Lazy import
AutoModel.registry["t5"] = T5Model                     # Direct

# Save and auto-load
config = BertConfig(hidden_size=1024)
model = BertModel(config=config)
model.save_pretrained("./model")
loaded = AutoModel.from_pretrained("./model")  # Auto-detects type

class Tokenizer(PretrainedMixin):
    def __init__(self, *args, vocab: dict[str, int] | None = None, **kwargs):
        super().__init__(*args, **kwargs)
        self.vocab = vocab or {}

    def save_pretrained(self, path):
        super().save_pretrained(path)
        # Save additional state (vocabulary)
        Path(path).joinpath("vocab.txt").write_text(...)

    @classmethod
    def from_pretrained(cls, path):
        config = cls.config_class.model_validate_json(...)
        vocab = ...  # Load vocabulary
        return cls(config=config, vocab=vocab)

from lazyregistry import ImportString

# Create an import string
import_str = ImportString("json:dumps")

# Load the object when needed
func = import_str.load()
func({"key": "value"})  # '{"key": "value"}'

registry = Registry(name="plugins")

# Dict-style assignment (auto-converts strings to ImportString)
registry["key"] = "module:object"      # Lazy
registry["key2"] = actual_object       # Immediate

# Configure behavior
registry.eager_load = True
registry["key3"] = "module:object"     # Eager load

from lazyregistry import NAMESPACE

# Direct access to .registry for registration
NAMESPACE["models"]["bert"] = "transformers:BertModel"
NAMESPACE["models"]["gpt2"] = GPT2Model

# Access registered items
model = NAMESPACE["models"]["bert"]

from lazyregistry.registry import LazyImportDict

registry = LazyImportDict()

# Configure behavior via attributes
registry.auto_import_strings = True   # Auto-convert strings to ImportString (default: True)
registry.eager_load = False           # Load immediately on assignment (default: False)

# Use like a normal dict
registry["key"] = "module:object"
registry.update({"key2": "module:object2"})

class MyConfig(PretrainedConfig):
    model_type: str = "my_model"

class MyModel(PretrainedMixin):
    config_class = MyConfig

model.save_pretrained("./path")
loaded = MyModel.from_pretrained("./path")

from lazyregistry.pretrained import PretrainedConfig, PretrainedMixin

# Each model has its own config class with type identifier
class BertConfig(PretrainedConfig):
    model_type: str = "bert"
    hidden_size: int = 768

class GPT2Config(PretrainedConfig):
    model_type: str = "gpt2"
    hidden_size: int = 768

# Base model class
class BaseModel(PretrainedMixin):
    config_class = PretrainedConfig

class AutoModel(AutoRegistry):
    registry = NAMESPACE["models"]
    config_class = PretrainedConfig
    type_key = "model_type"  # Can use any field name

# 1. Decorator registration - models inherit from BaseModel
@AutoModel.register_module("bert")
class BertModel(BaseModel):
    config_class = BertConfig

# 2. Direct registration via .registry
AutoModel.registry["gpt2"] = GPT2Model                   # Direct instance
AutoModel.registry["t5"] = "transformers:T5Model"        # Lazy import string

# 3. Bulk registration via .registry.update() - useful for many models
AutoModel.registry.update({
    "roberta": RobertaModel,
    "albert": "transformers:AlbertModel",
    "electra": "transformers:ElectraModel",
})

# Auto-detect and load
loaded = AutoModel.from_pretrained("./path")  # Detects type from config

# All imports happen upfront
from heavy_module_1 import ClassA
from heavy_module_2 import ClassB
from heavy_module_3 import ClassC

REGISTRY = {"a": ClassA, "b": ClassB, "c": ClassC}

# Import only what you use
from lazyregistry import Registry

registry = Registry(name="components")
registry.register("a", "heavy_module_1:ClassA")
registry.register("b", "heavy_module_2:ClassB")
registry.register("c", "heavy_module_3:ClassC")

# Only ClassA is imported
component = registry["a"]

# mypackage/__init__.py
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    # Static analysis sees these imports (IDE autocomplete, mypy, pyright)
    from .auto import AutoModel as AutoModel
    from .bert import BertModel as BertModel
    from .gpt2 import GPT2Model as GPT2Model
else:
    # Runtime uses lazy-loader (nothing imported until accessed)
    import lazy_loader as lazy

    __getattr__, __dir__, __all__ = lazy.attach(
        __name__,
        submod_attrs={
            "auto": ["AutoModel"],
            "bert": ["BertModel"],
            "gpt2": ["GPT2Model"],
        },
    )

# mypackage/auto.py
from lazyregistry import Registry
from lazyregistry.pretrained import AutoRegistry

class AutoModel(AutoRegistry):
    registry = Registry(name="models")

# String references - actual imports deferred until registry access
AutoModel.registry.update({
    "bert": "mypackage.bert:BertModel",
    "gpt2": "mypackage.gpt2:GPT2Model",
})