#!/usr/bin/env python3
from __future__ import annotations

from collections.abc import Callable

from accelerate import cpu_offload


class _ExecutionDeviceHint:
    execution_device = "cuda"
    offload = False
    hooks: tuple = ()


def move_registered_buffers(model, device: str) -> int:
    """Repair registered buffers missed by quantized-model placement wrappers."""
    moved = 0
    for module in model.modules():
        for name, buffer in getattr(module, "_buffers", {}).items():
            if buffer is not None and str(buffer.device) != device:
                module._buffers[name] = buffer.to(device)
                moved += 1
    return moved


def install_h3_rope_runtime_alignment(transformer) -> None:
    rope = transformer.rope
    original_forward = rope.forward

    def aligned_forward(position_ids):
        inv_freq = rope.inv_freq
        if str(inv_freq.device) != str(position_ids.device):
            rope._buffers["inv_freq"] = inv_freq.to(position_ids.device)
        return original_forward(position_ids)

    rope.forward = aligned_forward


def install_manual_h3_stage_offload(
    encoder_step_cls,
    *,
    text_encoder,
    transformer,
    empty_cuda_cache: Callable[[], None],
    place_transformer: bool = True,
    sequential_text_encoder: bool = False,
) -> None:
    """Run conditioning on CUDA, then free it before placing the denoiser."""
    install_h3_rope_runtime_alignment(transformer)
    original_encode_prompt = encoder_step_cls.encode_prompt
    original_call = encoder_step_cls.__call__

    def encode_prompt_on_cuda(
        components,
        prompt,
        images=None,
        device=None,
        dtype=None,
    ):
        return original_encode_prompt(
            components,
            prompt,
            images,
            device="cuda",
            dtype=dtype,
        )

    def call_then_place_denoiser(step, components, state):
        result = original_call(step, components, state)
        if not sequential_text_encoder:
            text_encoder.to("cpu")
            text_encoder._hf_hook = _ExecutionDeviceHint()
        empty_cuda_cache()
        if place_transformer:
            transformer.to("cuda")
            move_registered_buffers(transformer, "cuda")
        return result

    encoder_step_cls.encode_prompt = staticmethod(encode_prompt_on_cuda)
    encoder_step_cls.__call__ = call_then_place_denoiser
    if sequential_text_encoder:
        cpu_offload(text_encoder, execution_device="cuda", offload_buffers=True)
    else:
        text_encoder.to("cuda")