"""
Memory management module for SeedVR2
Handles VRAM usage, cache management, and memory optimization

Extracted from: seedvr2.py (lines 373-405, 607-626, 1016-1044)
"""

import torch
from modules.seedvr.src.common.cache import Cache
from modules.seedvr.src.models.dit_v2.rope import RotaryEmbeddingBase


def preinitialize_rope_cache(runner) -> None:
    """
    🚀 Pre-initialize RoPE cache to avoid OOM at first launch

    Args:
        runner: The model runner containing DiT and VAE models
    """

    # Create dummy tensors to simulate common shapes
    # Format: [batch, channels, frames, height, width] for vid_shape
    # Format: [batch, seq_len] for txt_shape
    common_shapes = [
        # Common video resolutions
        (torch.tensor([[1, 3, 3]], dtype=torch.long), torch.tensor([[77]], dtype=torch.long)),    # 1 frame, 77 tokens
        (torch.tensor([[4, 3, 3]], dtype=torch.long), torch.tensor([[77]], dtype=torch.long)),    # 4 frames
        (torch.tensor([[5, 3, 3]], dtype=torch.long), torch.tensor([[77]], dtype=torch.long)),    # 5 frames (4n+1 format)
        (torch.tensor([[1, 4, 4]], dtype=torch.long), torch.tensor([[77]], dtype=torch.long)),    # Higher resolution
    ]

    # Create mock cache for pre-initialization

    temp_cache = Cache()

    # Access RoPE modules in DiT (recursive search)
    def find_rope_modules(module):
        rope_modules = []
        for name, child in module.named_modules():
            if hasattr(child, 'get_freqs') and callable(child.get_freqs):
                rope_modules.append((name, child))
        return rope_modules

    rope_modules = find_rope_modules(runner.dit)

    # Pre-calculate for each RoPE module found
    for _name, rope_module in rope_modules:
        # Temporarily move module to CPU if necessary
        original_device = next(rope_module.parameters()).device if list(rope_module.parameters()) else torch.device('cpu')
        rope_module.to('cpu')

        for vid_shape, txt_shape in common_shapes:
            cache_key = f"720pswin_by_size_bysize_{tuple(vid_shape[0].tolist())}_sd3.mmrope_freqs_3d"

            def compute_freqs():
                # Calculate with reduced dimensions to avoid OOM
                with torch.no_grad():
                    # Detect RoPE module type
                    module_type = type(rope_module).__name__ # noqa: B023
                    if module_type == 'NaRotaryEmbedding3d':
                        # NaRotaryEmbedding3d: only takes shape (vid_shape)
                        return rope_module.get_freqs(vid_shape.cpu()) # noqa: B023
                    else:
                        # Standard RoPE: takes vid_shape and txt_shape
                        return rope_module.get_freqs(vid_shape.cpu(), txt_shape.cpu()) # noqa: B023

            # Store in cache
            temp_cache(cache_key, compute_freqs)

        rope_module.to(original_device)

    # Copy temporary cache to runner cache
    if hasattr(runner, 'cache'):
        runner.cache.cache.update(temp_cache.cache)
    else:
        runner.cache = temp_cache


def clear_rope_cache(runner) -> None:
    """
    🧹 Clear RoPE cache to free VRAM

    Args:
        runner: The model runner containing the cache
    """
    if hasattr(runner, 'cache') and hasattr(runner.cache, 'cache'):
        runner.cache.clear()

    if hasattr(runner, 'dit'):
        for module in runner.dit.modules():
            if isinstance(module, RotaryEmbeddingBase):
                if hasattr(module.get_axial_freqs, 'cache_clear'):
                    module.get_axial_freqs.cache_clear()
            if hasattr(module, 'cache') and hasattr(module.cache, 'clear'):
                module.cache.clear()