"""Direct prompt-only H3 packed-token construction."""

import math

import torch
import torch.nn.functional as functional


FRAME_PER_TOKEN = (1, 4, 4, 4, 4)
FRAME_RESCALE = 5.0 / 3.0


def patchify_video(latent: torch.Tensor) -> torch.Tensor:
    batch, channels, frames, height, width = latent.shape
    if batch != 1:
        raise ValueError("H3 supports batch size one.")
    if height % 2 or width % 2:
        raise ValueError("H3 video latent dimensions must be divisible by two.")
    return latent.reshape(batch, channels, frames, 1, height // 2, 2, width // 2, 2).permute(0, 2, 4, 6, 1, 3, 5, 7).reshape(-1, channels * 4)


def pack_audio(latent: torch.Tensor) -> torch.Tensor:
    return latent[0].permute(1, 2, 0).reshape(-1, latent.shape[1]).transpose(0, 1).contiguous().transpose(0, 1)


def unpatchify_video(rows: torch.Tensor, frames: int, latent_height: int, latent_width: int) -> torch.Tensor:
    height, width = latent_height // 2, latent_width // 2
    x = rows.reshape(1, frames, height, width, 24, 1, 2, 2).permute(0, 4, 1, 5, 2, 6, 3, 7)
    return x.reshape(1, 24, frames, latent_height, latent_width)


def _axis(dim: int, area: float) -> torch.Tensor:
    ratio, count = dim / area, dim // 2
    return (torch.arange(count, dtype=torch.float64) * ratio / count + (1 - ratio) / 2) * 32


def _frame_positions(height: int, width: int) -> torch.Tensor:
    """(t ignored) area-normalized (h, w) grid of one latent frame's 2x2-patch rows."""
    area = math.sqrt(height * width)
    ys, xs = torch.meshgrid(_axis(height, area), _axis(width, area), indexing="ij")
    return torch.stack((ys.flatten(), xs.flatten()), dim=-1)  # [frame_rows, 2]


def _video_positions(frames: int, height: int, width: int, offset: float) -> torch.Tensor:
    area = math.sqrt(height * width)
    ys, xs = torch.meshgrid(_axis(height, area), _axis(width, area), indexing="ij")
    spatial = torch.stack((ys.flatten(), xs.flatten()), dim=-1)
    spans = torch.tensor([FRAME_RESCALE * FRAME_PER_TOKEN[index % 5] for index in range(frames)], dtype=torch.float64)
    times = offset + torch.cat((torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)))
    result = torch.empty(frames, spatial.shape[0], 3, dtype=torch.float64)
    result[:, :, 0], result[:, :, 1:] = times[:, None], spatial[None]
    return result.reshape(-1, 3)


def _cond_positions(frames: int, cond_t: float, height: int, width: int) -> torch.Tensor:
    spatial = _frame_positions(height, width)
    g = torch.empty(spatial.shape[0], 3, dtype=torch.float64)
    g[:, 0] = cond_t
    g[:, 1:] = spatial
    return g


def _video_t_spans(n: int) -> list[float]:
    return [FRAME_RESCALE * FRAME_PER_TOKEN[k % 5] for k in range(n)]


def _audio_positions(steps: int, offset: float, width: int, height: int) -> torch.Tensor:
    area = math.sqrt(height * width)
    x_axis = _axis(width, area)
    result = torch.zeros(steps * 2, 3, dtype=torch.float64)
    result[:, 0] = (offset + torch.arange(steps, dtype=torch.float64)).repeat(2)
    result[:steps, 2], result[steps:, 2] = x_axis[0], x_axis[-1]
    return result


class H3PromptPacker:
    """Build `[text | audio | video]` tokens for prompt-only H3 T2V."""

    def __init__(self, checkpoint):
        self.video_weight = checkpoint.tensor("video_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
        self.video_bias = checkpoint.tensor("video_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32)
        self.audio_weight = checkpoint.tensor("audio_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
        self.audio_bias = checkpoint.tensor("audio_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32)
        self.text_weight = checkpoint.tensor("condition_proj.weight", dtype=torch.bfloat16)
        self.text_bias = checkpoint.tensor("condition_proj.bias", dtype=torch.bfloat16)

    VISUAL_COND_TIMESTEP = 0.999

    def __call__(
        self,
        text: torch.Tensor,
        video: torch.Tensor,
        audio: torch.Tensor,
        sigma: float | torch.Tensor,
        model_timesteps: torch.Tensor | None = None,
        *,
        text_token_tags: torch.Tensor | None = None,
        cond_latents: list[torch.Tensor] | None = None,
        frame_count: int | None = None,
        seed: int = 0,
    ) -> tuple[torch.Tensor, torch.Tensor, list[tuple[int, int, int]], torch.Tensor, tuple[int, int, int], tuple[int, int, int]]:
        """Build ``[text | cond | audio | video]`` packed rows for (fl2va) H3.

        ``text`` is the refined text span (width 5376 when already refined, 5120
        for raw Qwen states); ``text_token_tags`` is the per-token DiT modality tag
        (1=text, 0=video over vision pads). ``cond_latents`` are normalized keyframe
        latents ``[1,24,1,H/16,W/16]`` spliced right after the text as non-denoised
        cond rows with their own near-1 timestep. Returns
        ``(hidden, times, segments, positions, video_seg, audio_seg)`` where
        ``segments`` rows are ``t_row*3 + modality_tag``.
        """
        if text.shape[-1] == 5120:
            text_rows = functional.linear(text[0].to(self.text_weight.dtype), self.text_weight, self.text_bias).to(torch.bfloat16)
        elif text.shape[-1] == 5376:
            text_rows = text[0].to(torch.bfloat16)
        else:
            raise ValueError("H3 text states must be Qwen 5120-wide or refined 5376-wide.")
        text_length = text_rows.shape[0]
        latent_t, latent_h, latent_w = video.shape[2], video.shape[-2], video.shape[-1]
        frame_rows = (latent_h // 2) * (latent_w // 2)

        video_rows = functional.linear(patchify_video(video.to(torch.bfloat16)).float(), self.video_weight, self.video_bias).to(torch.bfloat16)
        audio_rows = functional.linear(pack_audio(audio.to(torch.bfloat16)).float(), self.audio_weight, self.audio_bias).to(torch.bfloat16)
        cond_rows = None
        if cond_latents:
            cond_patches = []
            # every cond video restarts the same CPU RNG stream (Comfy _cond_video_rows)
            for idx, z in enumerate(cond_latents):
                r = patchify_video(z.to(torch.float32))
                if self.VISUAL_COND_TIMESTEP < 1.0:
                    gen = torch.Generator("cpu").manual_seed(int(seed))
                    noise = torch.randn(r.shape, generator=gen, dtype=torch.float32)
                    r = self.VISUAL_COND_TIMESTEP * r + (1.0 - self.VISUAL_COND_TIMESTEP) * noise.to(r.device)
                cond_patches.append(r)
            cond_rows = functional.linear(torch.cat(cond_patches, dim=0), self.video_weight.to(torch.float32), self.video_bias.to(torch.float32)).to(torch.bfloat16)

        if model_timesteps is None:
            video_sigma = torch.as_tensor(sigma, device=text_rows.device, dtype=torch.float32).clamp(min=1e-6)
            base = video_sigma / (12.0 + video_sigma * (1.0 - 12.0))
            audio_sigma = 3.0 * base / (1.0 + (3.0 - 1.0) * base)
            video_time, audio_time = (1.0 - video_sigma).item(), (1.0 - audio_sigma).item()
        else:
            times_override = model_timesteps.to(device=text_rows.device, dtype=torch.float32).flatten()
            unique_override = sorted(times_override.tolist())
            video_time, audio_time = unique_override[0], unique_override[-1]

        has_vis_cond = cond_rows is not None
        cond_time = max(video_time, self.VISUAL_COND_TIMESTEP)
        unique_times = sorted({video_time, audio_time} | ({cond_time} if has_vis_cond else set()))
        t_row = {value: index for index, value in enumerate(unique_times)}
        times = torch.tensor(unique_times, device=text_rows.device, dtype=torch.float32)

        # Assemble hidden in segment order: text | cond | audio | video.
        parts = [text_rows]
        offsets = [0]
        for block in ((cond_rows, "cond"), (audio_rows, "audio"), (video_rows, "video")):
            if block[0] is None:
                continue
            parts.append(block[0])
            offsets.append(offsets[-1] + block[0].shape[0])
        hidden = torch.cat(parts, dim=0)

        audio_length = audio_rows.shape[0]
        cond_length = cond_rows.shape[0] if cond_rows is not None else 0

        # Positions: text rows, cond (first/last t anchors), audio, video.
        text_positions = torch.stack((torch.arange(text_length, dtype=torch.float64), torch.zeros(text_length), torch.zeros(text_length)), dim=-1)
        position_blocks = [text_positions]
        if cond_rows is not None and cond_latents:
            spans = _video_t_spans(latent_t)
            cond_t_values = [
                float(text_length) if idx == 0 else (float(text_length) + sum(spans) - FRAME_RESCALE if frame_count is not None and idx == len(cond_latents) - 1 else float(text_length))
                for idx in range(len(cond_latents))
            ]
            position_blocks.append(torch.cat([_cond_positions(frame_rows, cond_t, latent_h, latent_w) for cond_t in cond_t_values], dim=0))
        position_blocks.append(_audio_positions(audio.shape[-1], float(text_length + cond_length), latent_w, latent_h))
        position_blocks.append(_video_positions(latent_t, latent_h, latent_w, float(text_length + cond_length + audio_length)))
        positions = torch.cat(position_blocks, dim=0)

        # mod_segments: (start, stop, t_row*3 + tag).
        segments: list[tuple[int, int, int]] = []
        if text_token_tags is not None:
            tags = text_token_tags.view(-1).tolist()
            row_base = t_row[video_time] * 3
            run_start = 0
            for i in range(1, text_length + 1):
                if i == text_length or tags[i] != tags[run_start]:
                    segments.append((run_start, i, row_base + int(tags[run_start])))
                    run_start = i
        else:
            segments.append((0, text_length, t_row[video_time] * 3 + 1))

        cursor_start = text_length
        if cond_rows is not None:
            segments.append((cursor_start, cursor_start + cond_rows.shape[0], t_row[cond_time] * 3 + 0))
            cursor_start += cond_rows.shape[0]
        segments.append((cursor_start, cursor_start + audio_length, t_row[audio_time] * 3 + 2))
        cursor_start += audio_length
        video_start = cursor_start
        video_stop = video_start + video_rows.shape[0]
        segments.append((video_start, video_stop, t_row[video_time] * 3 + 0))

        video_segment = (video_start, video_stop, t_row[video_time])
        audio_start = text_length + cond_length
        audio_segment = (audio_start, audio_start + audio_length, t_row[audio_time])
        return hidden, times, segments, positions, video_segment, audio_segment