# Copyright (c) 2026 SandAI. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import math
from dataclasses import dataclass
from enum import IntEnum
from typing import Any, Literal, Optional, TYPE_CHECKING

import torch
from einops import rearrange
from inference.common import DataProxyConfig, Modality, VarlenHandler
from inference.model.dit.dit_module import FFAHandler
from torch.nn import functional as F
from unfoldNd import UnfoldNd

if TYPE_CHECKING:
    from inference.pipeline.video_generate import EvalInput


def calc_local_qk_range(num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field):
    token_per_frame = num_video_tokens // num_frames
    total_tokens = num_video_tokens + num_audio_and_txt_tokens

    q_range_list = []
    k_range_list = []

    for i in range(num_frames):
        local_q_range = torch.tensor([i * token_per_frame, (i + 1) * token_per_frame])
        local_k_range = torch.tensor(
            [(i - frame_receptive_field) * token_per_frame, (i + frame_receptive_field + 1) * token_per_frame]
        )

        q_range_list.append(local_q_range)
        k_range_list.append(local_k_range)
    local_q_range = torch.stack(q_range_list, dim=0)
    local_k_range = torch.stack(k_range_list, dim=0)

    local_k_range[local_k_range < 0] = 0
    local_k_range[local_k_range > num_video_tokens] = num_video_tokens

    video_q_range = torch.tensor([[0, num_video_tokens]])
    video_k_range = torch.tensor([[num_video_tokens, num_video_tokens + num_audio_and_txt_tokens]])

    at_q_ranges = torch.tensor([[num_video_tokens, total_tokens]])
    at_k_ranges = torch.tensor([[0, total_tokens]])

    q_ranges = torch.cat([local_q_range, video_q_range, at_q_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True)
    k_ranges = torch.cat([local_k_range, video_k_range, at_k_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True)

    return (q_ranges, k_ranges)


def calc_local_attn_ffa_handler(num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field):
    q_ranges, k_ranges = calc_local_qk_range(num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field)
    max_seqlen_q = num_video_tokens + num_audio_and_txt_tokens
    max_seqlen_k = num_video_tokens + num_audio_and_txt_tokens
    attn_type_map = torch.zeros([q_ranges.shape[0]], device="cuda", dtype=torch.int32)
    softmax_scale = None

    ffa_handler = FFAHandler(
        q_ranges=q_ranges,
        k_ranges=k_ranges,
        max_seqlen_q=max_seqlen_q,
        max_seqlen_k=max_seqlen_k,
        attn_type_map=attn_type_map,
        softmax_scale=softmax_scale,
    )
    return ffa_handler


def get_coords(
    shape: list[int],
    ref_feat_shape: list[int],
    offset_thw: list[int] = [0, 0, 0],
    device: torch.device = torch.device("cpu"),
    dtype: torch.dtype = torch.float32,
):
    """
    Generate feature-grid coordinates and corresponding original/reference size metadata.
    Args:
        feat_shape: [T, H, W] original feature-map shape
        ref_feat_shape: [T_ref, H_ref, W_ref] reference feature-map shape
        device: device for coordinate tensors
    Returns:
        coords: tensor shape (T*H*W, 9), containing (t, h, w, T, H, W, ref_T, ref_H, ref_W)
    """
    ori_t, ori_h, ori_w = shape
    ref_t, ref_h, ref_w = ref_feat_shape

    # Generate index ranges
    offset_t, offset_h, offset_w = offset_thw
    time_rng = torch.arange(ori_t, device=device, dtype=dtype) + offset_t
    height_rng = torch.arange(ori_h, device=device, dtype=dtype) + offset_h
    width_rng = torch.arange(ori_w, device=device, dtype=dtype) + offset_w

    # Use meshgrid to generate a 3D grid (T, H, W)
    time_grid, height_grid, width_grid = torch.meshgrid(time_rng, height_rng, width_rng, indexing="ij")

    # Stack and flatten
    coords_grid = torch.stack([time_grid, height_grid, width_grid], dim=-1)
    coords_flat = coords_grid.reshape(-1, 3)

    # Build and expand size metadata
    meta = torch.tensor([ori_t, ori_h, ori_w, ref_t, ref_h, ref_w], device=device, dtype=dtype)
    meta_expanded = meta.expand(coords_flat.size(0), -1)

    # Merge and return
    return torch.cat([coords_flat, meta_expanded], dim=-1)


@dataclass
class SingleData:
    video_x_t: torch.Tensor
    audio_x_t: torch.Tensor
    audio_feat_len: int
    txt_feat: torch.Tensor
    txt_feat_len: int
    t: int
    h: int
    w: int
    patch_size: int
    t_patch_size: int
    spatial_rope_interpolation: Literal["inter", "extra"]
    ref_audio_offset: int
    text_offset: int
    coords_style: Literal["v1", "v2"] = "v1"

    def __post_init__(self):
        self.video_token_num = self.video_x_t.shape[0]

        self.audio_x_t = self.audio_x_t[: self.audio_feat_len]
        self.txt_feat = self.txt_feat[: self.txt_feat_len]

        self.video_channel = self.video_x_t.shape[-1]
        self.audio_channel = self.audio_x_t.shape[-1]
        self.txt_channel = self.txt_feat.shape[-1]

    @property
    def device(self):
        return self.video_x_t.device

    @property
    def default_dtype(self):
        return self.video_x_t.dtype

    @property
    def total_token_num(self):
        return self.video_token_num + self.audio_feat_len + self.txt_feat_len

    @property
    def token_sequence(self):
        tensors_to_concat = [self.video_x_t, self.audio_x_t, self.txt_feat]
        max_channel = max(tensor.shape[-1] for tensor in tensors_to_concat)

        padded_tensors = [F.pad(t, (0, max_channel - t.shape[-1])) for t in tensors_to_concat]
        ret_val = torch.cat(padded_tensors, dim=0)
        return ret_val

    @property
    def modality_mapping(self):
        v_map = torch.full((self.video_token_num,), Modality.VIDEO, dtype=torch.int64, device=self.device)
        a_map = torch.full((self.audio_feat_len,), Modality.AUDIO, dtype=torch.int64, device=self.device)
        t_map = torch.full((self.txt_feat_len,), Modality.TEXT, dtype=torch.int64, device=self.device)

        modality_mapping = torch.cat([v_map, a_map, t_map], dim=0)
        return modality_mapping

    def default_coords(self, shape, ref_feat_shape, offset_thw=[0, 0, 0]):
        return get_coords(
            shape=shape, ref_feat_shape=ref_feat_shape, offset_thw=offset_thw, device=self.device, dtype=self.default_dtype
        )

    @property
    def coords_mapping(self):
        if self.spatial_rope_interpolation == "inter":
            video_ref_feat_shape = (self.t // self.t_patch_size, 32, 32)
        else:
            video_ref_feat_shape = (self.t // self.t_patch_size, self.h // self.patch_size, self.w // self.patch_size)

        video_coords = self.default_coords(
            shape=(self.t // self.t_patch_size, self.h // self.patch_size, self.w // self.patch_size),
            ref_feat_shape=video_ref_feat_shape,
        )

        if self.coords_style == "v1":
            audio_coords = self.default_coords(
                shape=(self.audio_feat_len, 1, 1), ref_feat_shape=(self.t // self.t_patch_size, 1, 1)
            )

            text_coords = self.default_coords(
                shape=(self.txt_feat_len, 1, 1), ref_feat_shape=(2, 1, 1), offset_thw=[self.text_offset, 0, 0]
            )

        elif self.coords_style == "v2":
            magic_audio_ref_t = (self.audio_feat_len - 1) // 4 + 1
            audio_coords = self.default_coords(
                shape=(self.audio_feat_len, 1, 1), ref_feat_shape=(magic_audio_ref_t // self.t_patch_size, 1, 1)
            )

            text_coords = self.default_coords(
                shape=(self.txt_feat_len, 1, 1), ref_feat_shape=(1, 1, 1), offset_thw=[-self.txt_feat_len, 0, 0]
            )

        coords_mapping = torch.cat([video_coords, audio_coords, text_coords], dim=0)
        return coords_mapping

    def depack_token_sequence(self, token_sequence):
        video_x_t = token_sequence[: self.video_token_num, : self.video_channel]
        video_x_t = rearrange(
            video_x_t,
            "(T H W) (pT pH pW C) -> C (T pT) (H pH) (W pW)",
            H=self.h // self.patch_size,
            W=self.w // self.patch_size,
            pT=self.t_patch_size,
            pH=self.patch_size,
            pW=self.patch_size,
        ).contiguous()

        audio_x_t = token_sequence[self.video_token_num : self.video_token_num + self.audio_feat_len, : self.audio_channel]
        return video_x_t, audio_x_t


@dataclass
class SimplePackedData:
    items: list[SingleData]

    @property
    def token_sequence(self):
        return torch.cat([item.token_sequence for item in self.items], dim=0)

    @property
    def modality_mapping(self):
        return torch.cat([item.modality_mapping for item in self.items], dim=0)

    @property
    def coords_mapping(self):
        return torch.cat([item.coords_mapping for item in self.items], dim=0)

    @property
    def total_token_num(self):
        return sum([item.total_token_num for item in self.items])

    def __getitem__(self, index):
        return self.items[index]

    @property
    def cu_seqlen(self):
        cu_seqlen = torch.cumsum(torch.tensor([item.total_token_num for item in self.items]), dim=0)
        cu_seqlen = torch.nn.functional.pad(cu_seqlen, (1, 0))
        return cu_seqlen

    @property
    def max_seqlen(self):
        return torch.tensor(max([item.total_token_num for item in self.items]))

    def depack_token_sequence(self, token_sequence):
        video_x_t_list = []
        audio_x_t_list = []

        token_sequence_list = torch.split(token_sequence, [item.total_token_num for item in self.items], dim=0)
        for item, token_sequence in zip(self.items, token_sequence_list):
            video_x_t, audio_x_t = item.depack_token_sequence(token_sequence)
            video_x_t_list.append(video_x_t)
            audio_x_t_list.append(audio_x_t)
        return torch.stack(video_x_t_list, dim=0), torch.stack(audio_x_t_list, dim=0)


class MagiDataProxy:
    def __init__(self, config: DataProxyConfig):
        self.patch_size = config.patch_size
        self.t_patch_size = config.t_patch_size
        self.frame_receptive_field = config.frame_receptive_field
        self.spatial_rope_interpolation = 'extra'
        self.ref_audio_offset = config.ref_audio_offset
        self.text_offset = config.text_offset
        self.unfold = UnfoldNd(
            kernel_size=(self.t_patch_size, self.patch_size, self.patch_size),
            stride=(self.t_patch_size, self.patch_size, self.patch_size),
        )
        self.coords_style = config.coords_style

        self._saved_data: dict[str, Any] = {}

    def saved_for_output(self, **kwargs):
        """
        Store intermediate data used by process_output.
        Supports keyword-argument style calls: saved_for_output(a=1, b=2)
        Can be called multiple times to accumulate data

        Args:
            **kwargs: key-value pairs to store
        """
        # Directly update dict; supports accumulation across calls
        self._saved_data.update(kwargs)

    def get_saved_data(self, key: str):
        """
        Get stored data
        """
        return self._saved_data[key]

    def img2tokens(self, x_t: torch.Tensor):
        x_t_unfolded = self.unfold(x_t)
        # Transpose dimensions from (N, col_dim, num_tokens) -> (N, num_tokens, col_dim)
        x_t = rearrange(x_t_unfolded, "N col_dim num_tokens -> N num_tokens col_dim").contiguous()
        return x_t

    def process_input(self, transported_data: "EvalInput"):
        # init img2col module

        batch_size, _, t, h, w = transported_data.x_t.shape
        # 1. Process video features while keeping the batch dimension
        x_t = self.img2tokens(transported_data.x_t)

        # 2. Process audio features while keeping the batch dimension
        # Assume transported_data.audio_x_t shape is already (N, num_tokens, col_dim)
        audio_x_t = transported_data.audio_x_t.contiguous()

        # Here we assume text_in shape is (N, num_tokens, col_dim)
        text_in = transported_data.txt_feat.contiguous()

        simple_packed_data = SimplePackedData(items=[])
        for i in range(batch_size):
            single_data = SingleData(
                video_x_t=x_t[i],
                audio_x_t=audio_x_t[i],
                audio_feat_len=transported_data.audio_feat_len[i],
                txt_feat=text_in[i],
                txt_feat_len=transported_data.txt_feat_len[i],
                t=t,
                h=h,
                w=w,
                patch_size=self.patch_size,
                t_patch_size=self.t_patch_size,
                spatial_rope_interpolation=self.spatial_rope_interpolation,
                ref_audio_offset=self.ref_audio_offset,
                text_offset=self.text_offset,
                coords_style=self.coords_style,
            )
            simple_packed_data.items.append(single_data)

        if self.frame_receptive_field != -1:
            assert batch_size == 1, "local attention only supports batch size 1"

            local_attn_handler = calc_local_attn_ffa_handler(
                num_video_tokens=simple_packed_data[0].video_token_num,
                num_audio_and_txt_tokens=simple_packed_data[0].audio_feat_len + simple_packed_data[0].txt_feat_len,
                num_frames=t,
                frame_receptive_field=self.frame_receptive_field,
            )
            if isinstance(local_attn_handler.max_seqlen_k, torch.Tensor):
                local_attn_handler.max_seqlen_k = local_attn_handler.max_seqlen_k.item()
            if isinstance(local_attn_handler.max_seqlen_q, torch.Tensor):
                local_attn_handler.max_seqlen_q = local_attn_handler.max_seqlen_q.item()
        else:
            local_attn_handler = None

        varlen_handler = VarlenHandler(
            cu_seqlens_q=simple_packed_data.cu_seqlen.to(torch.int32).cuda(),
            cu_seqlens_k=simple_packed_data.cu_seqlen.to(torch.int32).cuda(),
            max_seqlen_q=simple_packed_data.max_seqlen.to(torch.int32).cuda(),
            max_seqlen_k=simple_packed_data.max_seqlen.to(torch.int32).cuda(),
        )

        self.saved_for_output(simple_packed_data=simple_packed_data)

        x = simple_packed_data.token_sequence
        coords_mapping = simple_packed_data.coords_mapping
        modality_mapping = simple_packed_data.modality_mapping

        return (x, coords_mapping, modality_mapping, varlen_handler, local_attn_handler)

    def process_output(self, x: torch.Tensor):
        # Inserting operations in between may corrupt parallel-runtime data and cause latent errors

        simple_packed_data: SimplePackedData = self.get_saved_data("simple_packed_data")
        x_video, x_audio = simple_packed_data.depack_token_sequence(x)

        return (x_video, x_audio)