"""Compatibility helper for loading AOTI-compiled repeated transformer blocks."""

from typing import cast

import torch
from huggingface_hub import hf_hub_download
from spaces.zero.torch.aoti import ZeroGPUCompiledModel, ZeroGPUWeights
from torch._functorch._aot_autograd.subclass_parametrization import (
    unwrap_tensor_subclass_parameters,
)


def _shallow_clone_module(module: torch.nn.Module) -> torch.nn.Module:
    clone = object.__new__(module.__class__)
    clone.__dict__ = module.__dict__.copy()
    clone._parameters = module._parameters.copy()
    clone._buffers = module._buffers.copy()
    clone._modules = {
        name: _shallow_clone_module(child)
        for name, child in module._modules.items()
        if child is not None
    }
    return clone


def aoti_blocks_load(
    module: torch.nn.Module,
    repo_id: str,
    variant: str | None = None,
) -> None:
    repeated_blocks = cast(list[str], module._repeated_blocks)
    aoti_files = {
        name: hf_hub_download(
            repo_id=repo_id,
            filename="package.pt2",
            subfolder=name if variant is None else f"{name}.{variant}",
        )
        for name in repeated_blocks
    }

    for block_name, aoti_file in aoti_files.items():
        for block in module.modules():
            if block.__class__.__name__ != block_name:
                continue
            cloned_block = _shallow_clone_module(block)
            unwrap_tensor_subclass_parameters(cloned_block)
            weights = ZeroGPUWeights(cloned_block.state_dict())
            block.forward = ZeroGPUCompiledModel(aoti_file, weights)