# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

from __future__ import annotations

import math
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from typing import Any

import torch
from vllm.logger import init_logger

from vllm_omni.diffusion.cache.seacache.config import SeaCacheConfig
from vllm_omni.diffusion.cache.seacache.sea_filter import (
    apply_sea_filter,
    extrapolate_residual,
    indicator_distance,
)
from vllm_omni.diffusion.cache.seacache.state import SeaCacheState
from vllm_omni.diffusion.cache.teacache.extractors import CacheContext, get_extractor
from vllm_omni.diffusion.hooks import HookRegistry, ModelHook, StateManager

logger = init_logger(__name__)


def _is_parameter_sharded(module: torch.nn.Module) -> bool:
    """Detect parameter-sharding runtimes whose collectives cannot be skipped."""
    for submodule in module.modules():
        module_type = type(submodule)
        if callable(getattr(submodule, "_get_fsdp_state", None)):
            return True
        if module_type.__name__ == "FullyShardedDataParallel" and module_type.__module__.startswith(
            "torch.distributed.fsdp"
        ):
            return True
        for parameter in submodule.parameters(recurse=False):
            parameter_type = type(parameter)
            if (
                parameter_type.__name__ == "FlatParameter"
                and parameter_type.__module__.startswith("torch.distributed.fsdp")
            ) or (
                parameter_type.__name__ == "DTensor"
                and parameter_type.__module__.startswith("torch.distributed.tensor")
            ):
                return True
    return False


class SeaCacheRootHook(ModelHook):
    """Drive SeaCache gating and transformer forward control."""

    _HOOK_NAME = "sea_cache"

    def __init__(
        self,
        config: SeaCacheConfig,
        *,
        current_step_callback: Callable[[], int | torch.Tensor | None] | None = None,
        current_sigma_callback: Callable[[], float | torch.Tensor | None] | None = None,
        num_inference_steps_callback: Callable[[], int | torch.Tensor | None] | None = None,
        extractor_fn: Callable[..., CacheContext] | None = None,
    ) -> None:
        super().__init__()
        self.config = config
        self.current_step_callback = current_step_callback
        self.current_sigma_callback = current_sigma_callback
        self.num_inference_steps_callback = num_inference_steps_callback
        self.state_manager = StateManager(SeaCacheState)
        self._warned_messages: set[str] = set()
        self.full_count = 0
        self.skip_count = 0
        self.extractor_fn = extractor_fn
        self._parameter_sharded = False
        self._collective_skip_groups: list[torch.distributed.ProcessGroup] = []

    def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
        if self.extractor_fn is None:
            self.extractor_fn = get_extractor(type(module))
        self._parameter_sharded = _is_parameter_sharded(module)
        seen_groups: set[int] = set()
        for block in getattr(module, "gen_layers", ()):
            registry = getattr(block, "_hook_registry", None)
            dlo_hook = registry.get_hook("distributed_layerwise_offload") if registry is not None else None
            group = getattr(dlo_hook, "dp_group", None)
            if group is not None and int(getattr(dlo_hook, "dp_size", 1)) > 1 and id(group) not in seen_groups:
                seen_groups.add(id(group))
                self._collective_skip_groups.append(group)
        return module

    def _warn_once(self, message: str) -> None:
        if message not in self._warned_messages:
            logger.warning(message)
            self._warned_messages.add(message)

    @contextmanager
    def cache_context(self, name: str) -> Iterator[None]:
        previous_context = self.state_manager._context
        self.state_manager.set_context(name)
        try:
            yield
        finally:
            self.state_manager.set_context(previous_context)

    def _build_indicator(
        self,
        vision_items: list[torch.Tensor] | None,
        sigma: float,
    ) -> list[torch.Tensor] | None:
        if not vision_items:
            return None
        hidden_states = vision_items[-1]
        if not isinstance(hidden_states, torch.Tensor) or hidden_states.ndim != 5:
            return None

        # Controls precede the denoised target in packed vision-token order,
        # and each vision item is filtered independently.
        if any(
            not isinstance(item, torch.Tensor)
            or item.ndim != 5
            or item.shape[0] != hidden_states.shape[0]
            or item.shape[1:] != hidden_states.shape[1:]
            for item in vision_items
        ):
            return None

        indicator = []
        for batch_index in range(hidden_states.shape[0]):
            for latent in vision_items:
                thwc = latent[batch_index].movedim(0, -1)
                indicator.append(
                    apply_sea_filter(
                        thwc,
                        sigma=sigma,
                        power_exp=self.config.power_exp,
                    ).detach()
                )
        return indicator or None

    def _resolve_gate(
        self,
        state: SeaCacheState,
        indicator: list[torch.Tensor] | None,
        step: int,
        num_inference_steps: int,
    ) -> bool:
        if state.last_step is not None and step != state.last_step + 1:
            state.reset()
        state.last_step = step
        max_consecutive = bool(
            self.config.max_consecutive_cached and state.consecutive_cached >= self.config.max_consecutive_cached
        )
        forced_compute = (
            step < 1
            or step >= num_inference_steps - 1
            or max_consecutive
            or not state.history
            or indicator is None
            or state.previous_indicator is None
        )
        if forced_compute:
            state.accumulated_distance = 0.0
            state.previous_indicator = None if indicator is None else [value.detach() for value in indicator]
            return True

        assert indicator is not None
        assert state.previous_indicator is not None
        distance = indicator_distance(indicator, state.previous_indicator)
        state.previous_indicator = [value.detach() for value in indicator]
        if not math.isfinite(distance):
            state.accumulated_distance = 0.0
            self._warn_once("SeaCache indicator history changed shape, device, or dtype; running full.")
            return True

        state.accumulated_distance += distance
        if state.accumulated_distance < self.config.threshold:
            return False
        state.accumulated_distance = 0.0
        return True

    def _synchronize_compute(self, compute: bool, device: torch.device) -> bool:
        if not torch.distributed.is_available() or not torch.distributed.is_initialized():
            return True if self._parameter_sharded else compute
        decision = torch.tensor(int(compute), dtype=torch.int32, device=device)
        if self._parameter_sharded:
            from vllm_omni.diffusion.distributed.parallel_state import (
                get_fs_group,
                get_sequence_parallel_world_size,
                get_sp_group,
            )

            fs_group = get_fs_group()
            if fs_group.world_size > 1:
                torch.distributed.all_reduce(
                    decision,
                    op=torch.distributed.ReduceOp.MAX,
                    group=fs_group.device_group,
                )
            if get_sequence_parallel_world_size() > 1:
                torch.distributed.all_reduce(
                    decision,
                    op=torch.distributed.ReduceOp.MAX,
                    group=get_sp_group().device_group,
                )
            return bool(decision.item())

        for group in self._collective_skip_groups:
            torch.distributed.all_reduce(
                decision,
                op=torch.distributed.ReduceOp.MAX,
                group=group,
            )
        from vllm_omni.diffusion.distributed.parallel_state import (
            get_sequence_parallel_world_size,
            get_sp_group,
        )

        if get_sequence_parallel_world_size() > 1:
            torch.distributed.all_reduce(
                decision,
                op=torch.distributed.ReduceOp.MAX,
                group=get_sp_group().device_group,
            )
        return bool(decision.item())

    @torch.compiler.disable
    def new_forward(
        self,
        module: torch.nn.Module,
        *args: Any,
        **kwargs: Any,
    ) -> Any:
        if self.extractor_fn is None:
            raise RuntimeError("SeaCache extractor was not initialized")
        ctx = self.extractor_fn(module, *args, **kwargs)

        if torch.is_grad_enabled():
            self._warn_once("SeaCache is inference-only; autograd-enabled calls run in full.")
            return self._run_uncached(ctx)
        if self.state_manager._current_context is None:
            self._warn_once("SeaCache requires an explicit cache context; running full.")
            return self._run_uncached(ctx)
        callbacks = (
            self.current_step_callback,
            self.current_sigma_callback,
            self.num_inference_steps_callback,
        )
        if any(callback is None for callback in callbacks):
            self._warn_once("SeaCache requires scheduler step, sigma, and step-count callbacks; running full.")
            return self._run_uncached(ctx)
        assert self.current_step_callback is not None
        assert self.current_sigma_callback is not None
        assert self.num_inference_steps_callback is not None

        try:
            extra_states = ctx.extra_states or {}
            vision_items = extra_states.get("sea_cache_latents")
            if not isinstance(vision_items, list):
                raise ValueError("extractor did not provide SeaCache vision inputs")
            noisy_frame_mask = extra_states.get("sea_cache_noisy_frame_mask")
            conditioning_only = isinstance(noisy_frame_mask, torch.Tensor) and not bool(
                torch.any(noisy_frame_mask != 0).item()
            )

            step = self.current_step_callback()
            sigma = self.current_sigma_callback()
            num_inference_steps = self.num_inference_steps_callback()
            if isinstance(step, torch.Tensor):
                step = step.item()
            if isinstance(sigma, torch.Tensor):
                sigma = sigma.item()
            if isinstance(num_inference_steps, torch.Tensor):
                num_inference_steps = num_inference_steps.item()
            if step is None or sigma is None or num_inference_steps is None:
                raise ValueError("scheduler metadata is unavailable")
            step = int(step)
            sigma = float(sigma)
            num_inference_steps = int(num_inference_steps)
            if (
                step < 0
                or num_inference_steps <= 0
                or step >= num_inference_steps
                or not math.isfinite(sigma)
                or not 0.0 <= sigma <= 1.0
            ):
                raise ValueError("expected a valid step index and exact sigma in [0, 1]")
        except (IndexError, TypeError, ValueError, RuntimeError) as error:
            self._warn_once(f"SeaCache metadata is invalid; running full: {error}")
            return self._run_uncached(ctx)

        if conditioning_only:
            self._warn_once("SeaCache requires noisy vision; conditioning-only calls run in full.")
            return self._run_uncached(ctx)

        state: SeaCacheState = self.state_manager.get_state()
        try:
            indicator = self._build_indicator(vision_items, sigma)
        except (TypeError, ValueError, RuntimeError) as error:
            self._warn_once(f"SeaCache could not construct its vision indicator; running full: {error}")
            indicator = None

        local_compute = self._resolve_gate(state, indicator, step, num_inference_steps)
        should_compute = self._synchronize_compute(local_compute, ctx.hidden_states.device)
        if should_compute and not local_compute:
            state.accumulated_distance = 0.0

        if should_compute:
            self.full_count += 1
            output = self._run_full_stack(ctx)
            result = ctx.postprocess(output)
            self._record_execution(state, step, ctx.hidden_states, output)
            return result

        residual = extrapolate_residual(
            state.history,
            step,
            self.config.residual_order,
        )
        if residual.device != ctx.hidden_states.device:
            residual = residual.to(ctx.hidden_states.device)
        state.consecutive_cached += 1
        self.skip_count += 1

        can_reuse = (
            residual.shape == ctx.hidden_states.shape
            and residual.device == ctx.hidden_states.device
            and residual.dtype == ctx.hidden_states.dtype
        )
        if can_reuse:
            return ctx.postprocess(ctx.hidden_states + residual)

        output = self._run_full_stack(ctx)
        result = ctx.postprocess(output)
        self._record_execution(state, step, ctx.hidden_states, output)
        return result

    @staticmethod
    def _run_full_stack(ctx: CacheContext) -> torch.Tensor:
        outputs = ctx.run_transformer_blocks()
        if not outputs:
            raise RuntimeError("Cache extractor returned no transformer outputs")
        return outputs[0]

    @staticmethod
    def _run_uncached(ctx: CacheContext) -> Any:
        return ctx.postprocess(SeaCacheRootHook._run_full_stack(ctx))

    def _record_execution(
        self,
        state: SeaCacheState,
        step: int,
        execution_input: torch.Tensor,
        output: torch.Tensor,
    ) -> None:
        if (
            output.shape == execution_input.shape
            and output.device == execution_input.device
            and output.dtype == execution_input.dtype
        ):
            state.history.append((step, (output - execution_input).detach().clone()))
            state.history = state.history[-(self.config.residual_order + 1) :]
            state.consecutive_cached = 0
            return

        state.history.clear()
        state.accumulated_distance = 0.0
        self._warn_once("SeaCache execution boundary returned an incompatible tensor; clearing cache history.")

    def reset_state(self, module: torch.nn.Module) -> torch.nn.Module:
        self.state_manager.reset()
        self.full_count = 0
        self.skip_count = 0
        return module

    def refresh(self, module: torch.nn.Module) -> None:
        self.reset_state(module)


def apply_sea_cache_hook(
    module: torch.nn.Module,
    config: SeaCacheConfig,
    *,
    current_step_callback: Callable[[], int | torch.Tensor | None] | None = None,
    current_sigma_callback: Callable[[], float | torch.Tensor | None] | None = None,
    num_inference_steps_callback: Callable[[], int | torch.Tensor | None] | None = None,
    extractor_fn: Callable[..., CacheContext] | None = None,
) -> SeaCacheRootHook:
    registry = HookRegistry.get_or_create(module)
    hook = SeaCacheRootHook(
        config,
        current_step_callback=current_step_callback,
        current_sigma_callback=current_sigma_callback,
        num_inference_steps_callback=num_inference_steps_callback,
        extractor_fn=extractor_fn,
    )
    registry.register_hook(SeaCacheRootHook._HOOK_NAME, hook)
    return hook