"""Single-threaded port of `xvideo.serving.joyomni_streaming.JoyOmniV2VStreamingSession`.

Upstream runs the chunk pipeline across five worker threads (VAE encode -> DiT ->
VAE decode -> pseudo-latent encode -> JPEG post-process) feeding a WebSocket
downlink.  A Gradio Space has no downlink and ZeroGPU gives one GPU slice per
request, so the same stages are executed sequentially here and the decoded
frames are yielded chunk by chunk.  Every tensor-level step (chunk windows, KV
cache ids/eviction, the clean-KV store pass, the pseudo-latent decode trick,
RoPE temporal ids) mirrors the reference implementation.
"""

from __future__ import annotations

from contextlib import nullcontext
from dataclasses import dataclass
from typing import Any, Iterator

import numpy as np
import torch
from diffusers.utils.torch_utils import randn_tensor
from einops import rearrange
from PIL import Image

from xvideo.config import generate_video_image_bucket
from xvideo.models.pipeline import PRECISION_TO_TYPE
from xvideo.models.vae import vae_compile as _vc
from xvideo.utils import _dynamic_resize_from_bucket, seed_everything

DEFAULT_REFERENCE_IMG_IV2V_BASESIZE = 512


def _autocast_ctx(device_type: str, dtype: torch.dtype, enabled: bool):
    if device_type in {"cuda", "cpu"}:
        return torch.autocast(device_type=device_type, dtype=dtype, enabled=enabled)
    return nullcontext()


@dataclass(frozen=True)
class StreamingSettings:
    """Mirrors the upstream dataclass; defaults match `serve_joyomni_streaming.py`."""

    height: int = 720
    width: int = 1248
    num_inference_steps: int = 2
    seed: int = 42
    max_sequence_length: int | None = None
    enable_denormalization: bool | None = None
    max_temporal_ids: int | None = None
    store_clean_self_only: bool = True


class StreamingEditor:
    """Chunk-by-chunk video editing session over a live JoyAI-Video-Edit pipeline."""

    def __init__(self, cfg, pipeline, settings: StreamingSettings):
        self.cfg = cfg
        self.pipeline = pipeline
        self.settings = settings

        self.device = pipeline.transformer.device
        self.device_type = torch.device(self.device).type
        self.target_dtype = PRECISION_TO_TYPE[cfg.dit_precision]
        self.vae_dtype = PRECISION_TO_TYPE[cfg.vae_precision]
        self.autocast_enabled = self.target_dtype != torch.float32
        self.vae_autocast_enabled = self.vae_dtype != torch.float32

        self.chunk_size = pipeline._resolve_streaming_chunk_size(None, 1)
        if self.chunk_size != 1:
            raise ValueError("Streaming editing requires transformer chunk_size=1.")

        self.ffactor_t = int(pipeline.vae_scale_factor_temporal)
        self.latent_channels = int(pipeline.vae.config.latent_channels)

        stem = getattr(pipeline.vae, "stem", None)
        canon_h, canon_w = settings.height, settings.width
        if stem is not None and settings.height % stem.stride == 0 and settings.width % stem.stride == 0:
            canon_h = settings.height * stem.group // stem.stride
            canon_w = settings.width * stem.group // stem.stride
        self.latent_h = canon_h // int(pipeline.vae_scale_factor)
        self.latent_w = canon_w // int(pipeline.vae_scale_factor)

        self.local_window_size = int(getattr(pipeline.transformer.config, "local_window_size", 1))
        self.global_sink_chunk = pipeline._resolve_global_sink_chunk(None, pipeline.transformer)
        self.enable_denormalization = (
            cfg.enable_denormalization
            if settings.enable_denormalization is None
            else settings.enable_denormalization
        )

        self.generator = torch.Generator(device=self.device).manual_seed(settings.seed)
        self.chunk_idx = 0
        self.prev_source_frame: torch.Tensor | None = None
        self.pseudo_latent: torch.Tensor | None = None
        self.ref_image_latent: torch.Tensor | None = None
        self.ref_image_kv_prefilled = False
        self.streaming_cond_embeds: torch.Tensor | None = None
        self.streaming_cond_mask: torch.Tensor | None = None

    # ------------------------------------------------------------------ utils
    def _clear_vae_caches(self) -> None:
        clear = getattr(self.pipeline.vae, "clear_cache", None)
        if callable(clear):
            clear()

    def _resize_frame(self, frame: Image.Image) -> Image.Image:
        frame = frame.convert("RGB")
        if frame.size == (self.settings.width, self.settings.height):
            return frame
        resampling = getattr(Image, "Resampling", Image).BICUBIC
        return frame.resize((self.settings.width, self.settings.height), resampling)

    def _frames_to_tensor(self, frames: list[Image.Image]) -> torch.Tensor:
        arrays = [np.asarray(self._resize_frame(f), dtype=np.uint8) for f in frames]
        pixel = torch.from_numpy(np.stack(arrays, axis=0))
        pixel = rearrange(pixel, "t h w c -> 1 c t h w").to(torch.float32)
        return pixel / 127.5 - 1.0

    # -------------------------------------------------------------- init pass
    @torch.no_grad()
    def start(self, prompt: str, first_frame: Image.Image, ref_image: Image.Image | None = None) -> None:
        transformer = self.pipeline.transformer
        if not getattr(transformer.config, "causal", False):
            raise ValueError("Streaming editing requires a causal transformer config.")
        transformer.config.use_inference_kv_cache = True
        self.pipeline._interrupt = False

        seed_everything(self.settings.seed)
        self.prompt = prompt
        self.ref_image = ref_image.convert("RGB") if ref_image is not None else None
        self.ref_image_latent = self._encode_ref_image_latent()
        self._encode_streaming_prompt(prompt, first_frame)

        self._clear_vae_caches()
        transformer.reset_inference_kv_cache()
        if self.ref_image_latent is not None:
            self.pipeline._prefill_static_reference_kv_cache(
                transformer,
                prompt_embeds=self.streaming_cond_embeds,
                prompt_embeds_mask=self.streaming_cond_mask,
                reference_image_latents=self.ref_image_latent,
                transformer_dtype=self.target_dtype,
            )
            self.ref_image_kv_prefilled = True

    @torch.no_grad()
    def _encode_streaming_prompt(self, prompt: str, anchor_frame: Image.Image) -> None:
        prompt_image = self._resize_frame(anchor_frame)
        templated = f"<|im_start|>user\n<image>\n{prompt}<|im_end|>\n"
        max_sequence_length = self.settings.max_sequence_length or int(self.cfg.text_token_max_length)
        embeds, mask = self.pipeline.encode_prompt(
            prompt=[templated],
            images=[prompt_image],
            device=self.device,
            num_videos_per_prompt=1,
            max_sequence_length=max_sequence_length,
            template_type="video",
        )
        self.streaming_cond_embeds = embeds
        self.streaming_cond_mask = mask

    @torch.no_grad()
    def _encode_ref_image_latent(self) -> torch.Tensor | None:
        if self.ref_image is None:
            return None
        basesize = getattr(self.cfg, "ref_image_basesize", DEFAULT_REFERENCE_IMG_IV2V_BASESIZE)
        buckets = generate_video_image_bucket(
            img_basesize=basesize, bs_img=1, bs_vid=0, bs_mimg=0, bs_mvid=0,
        )
        resized, _ = _dynamic_resize_from_bucket(
            self.ref_image, bucket_configs=buckets, num_frames=1, num_items=1, return_bucket=True,
        )
        pixel = torch.from_numpy(np.asarray(resized))
        pixel = rearrange(pixel, "h w c -> c h w")
        if pixel.dtype != torch.uint8:
            pixel = pixel.clamp(0, 255).to(torch.uint8)
        normalized = pixel.to(torch.float32) / 127.5 - 1.0
        ref = rearrange(normalized, "c h w -> 1 c 1 h w").to(device=self.device, dtype=self.vae_dtype)

        self._clear_vae_caches()
        encoded = _vc.encode_via_dynamic(self.pipeline.vae, ref)
        latent = encoded.latent_dist.sample() if hasattr(encoded, "latent_dist") else encoded
        if self.enable_denormalization:
            latent = self.pipeline.normalize_latents(latent)
        self._clear_vae_caches()
        return latent[:, :, :1].to(device=self.device, dtype=self.target_dtype)

    # ----------------------------------------------------------- chunk stages
    @property
    def frames_per_next_chunk(self) -> int:
        return 1 if self.chunk_idx == 0 else self.ffactor_t

    @torch.no_grad()
    def _encode_reference_chunk(self, source_frames: list[Image.Image]) -> torch.Tensor:
        if self.chunk_idx == 0:
            source_window = self._frames_to_tensor(source_frames[:1])
        else:
            if self.prev_source_frame is None:
                raise RuntimeError("Missing previous source frame for streaming VAE encode.")
            if len(source_frames) != self.ffactor_t:
                raise ValueError(f"Expected {self.ffactor_t} frames, got {len(source_frames)}.")
            new_frames = self._frames_to_tensor(source_frames)
            source_window = torch.cat([self.prev_source_frame, new_frames], dim=2)

        self.prev_source_frame = source_window[:, :, -1:].detach().cpu()
        source_window = source_window.to(device=self.device, dtype=self.target_dtype)
        source_window = _vc.prep_input(source_window)

        with _autocast_ctx(self.device_type, self.vae_dtype, self.vae_autocast_enabled):
            ref_latent = self.pipeline._sample_vae_latents(
                source_window, enable_denormalization=self.enable_denormalization,
            )
        return ref_latent[:, :, -self.chunk_size:].to(device=self.device, dtype=self.target_dtype)

    @torch.no_grad()
    def _denoise_chunk(self, ref_chunk_latent: torch.Tensor) -> torch.Tensor:
        pipeline = self.pipeline
        chunk_idx = self.chunk_idx
        noise_shape = (1, self.latent_channels, self.chunk_size, self.latent_h, self.latent_w)
        current = randn_tensor(
            noise_shape, generator=self.generator, device=self.device, dtype=self.target_dtype,
        )

        total_latent_frames = chunk_idx + 1
        chunk_window = pipeline._get_chunk_windows(
            total_latent_frames=total_latent_frames,
            chunk_size=self.chunk_size,
            window_size=self.local_window_size,
            global_sink_chunk=self.global_sink_chunk,
        )[-1]
        selected_chunk_ids = chunk_window["selected_chunk_ids"]
        history_chunk_ids = selected_chunk_ids[:-1]
        active_chunk_id = selected_chunk_ids[-1]

        window_ids = pipeline._gather_window_temporal_ids(
            selected_chunk_ids,
            self.chunk_size,
            total_latent_frames,
            self.device,
            max_temporal_ids=self.settings.max_temporal_ids,
        )
        current_chunk_temporal_ids = window_ids[-self.chunk_size:]
        cached_temporal_ids = window_ids[: -self.chunk_size]
        if cached_temporal_ids.numel() == 0:
            cached_temporal_ids = None

        pipeline.scheduler.set_timesteps(self.settings.num_inference_steps, device=self.device)

        cache_memory_ids = [pipeline._kv_cache_memory_id("clean", cid) for cid in history_chunk_ids]
        if self.ref_image_kv_prefilled:
            cache_memory_ids.append(pipeline._kv_cache_memory_id("ref_image"))

        for timestep in pipeline.scheduler.timesteps:
            with _autocast_ctx(self.device_type, self.target_dtype, self.autocast_enabled):
                latent_model_input = current.to(self.target_dtype)
                t_expand = timestep.repeat(latent_model_input.shape[0])
                with pipeline.transformer.cache_context("cond"):
                    noise_pred = pipeline.transformer(
                        hidden_states=latent_model_input,
                        timestep=t_expand,
                        encoder_hidden_states=self.streaming_cond_embeds,
                        encoder_hidden_states_mask=self.streaming_cond_mask,
                        ref_video_latent=ref_chunk_latent,
                        current_temporal_ids=current_chunk_temporal_ids.unsqueeze(0).expand(
                            latent_model_input.shape[0], -1
                        ),
                        cached_temporal_ids=(
                            cached_temporal_ids.unsqueeze(0).expand(latent_model_input.shape[0], -1)
                            if cached_temporal_ids is not None
                            else None
                        ),
                        kv_cache_mode="reuse",
                        kv_cache_scope="cond",
                        kv_cache_chunk_id=active_chunk_id,
                        kv_cache_selected_chunk_ids=cache_memory_ids,
                        kv_cache_pre_rope=True,
                    )[0]
                current = pipeline.scheduler.step(
                    noise_pred, timestep, current.clone(), return_dict=False,
                )[0]

        keep_before_store = {pipeline._kv_cache_memory_id("clean", cid) for cid in history_chunk_ids}
        if self.ref_image_kv_prefilled:
            keep_before_store.add(pipeline._kv_cache_memory_id("ref_image"))
        pipeline.transformer.evict_kv_cache_chunks(keep_before_store)

        store_self_only = self.settings.store_clean_self_only
        pipeline._store_clean_chunk_kv_cache(
            pipeline.transformer,
            clean_chunk_latents=current.to(self.target_dtype),
            chunk_temporal_ids=current_chunk_temporal_ids.unsqueeze(0).expand(current.shape[0], -1),
            prompt_embeds=self.streaming_cond_embeds,
            prompt_embeds_mask=self.streaming_cond_mask,
            active_chunk_id=active_chunk_id,
            history_chunk_ids=[] if store_self_only else history_chunk_ids,
            pre_rope=True,
            cached_temporal_ids=None if store_self_only else cached_temporal_ids,
            store_mode="store" if store_self_only else "reuse_store",
        )
        return current

    @torch.no_grad()
    def _decode_chunk_pixels(self, latents: torch.Tensor) -> torch.Tensor:
        pipeline = self.pipeline
        vae = pipeline.vae
        flat = pipeline.denormalize_latents(latents) if self.enable_denormalization else latents

        if self.chunk_idx > 0:
            if self.pseudo_latent is None:
                raise RuntimeError("Missing pseudo latent for streaming decode.")
            decode_input = torch.cat([self.pseudo_latent, flat], dim=2)
            self.pseudo_latent = None
        else:
            decode_input = flat

        self._clear_vae_caches()
        decode_input = _vc.prep_input(decode_input)
        with _autocast_ctx(self.device_type, self.vae_dtype, self.vae_autocast_enabled):
            decoded = vae.decode(decode_input, return_dict=False)[0]
        self._clear_vae_caches()

        if self.chunk_idx > 0:
            decoded = decoded[:, :, -(self.chunk_size * self.ffactor_t):]
        return decoded.detach()

    @torch.no_grad()
    def _encode_pseudo_latent(self, decoded_pixels: torch.Tensor) -> None:
        prev = decoded_pixels[:, :, -1:].detach().to(device=self.device, dtype=self.vae_dtype)
        self._clear_vae_caches()
        prev = _vc.prep_input(prev)
        with _autocast_ctx(self.device_type, self.vae_dtype, self.vae_autocast_enabled):
            enc = self.pipeline.vae.encode(prev)
            latent = enc.latent_dist.sample() if hasattr(enc, "latent_dist") else enc
        self._clear_vae_caches()
        self.pseudo_latent = latent.detach()

    def _evict_after_store(self) -> None:
        pipeline = self.pipeline
        windows = pipeline._get_chunk_windows(
            total_latent_frames=self.chunk_idx + 2,
            chunk_size=self.chunk_size,
            window_size=self.local_window_size,
            global_sink_chunk=self.global_sink_chunk,
        )
        keep = {pipeline._kv_cache_memory_id("clean", cid) for cid in windows[-1]["selected_chunk_ids"]}
        if self.ref_image_kv_prefilled:
            keep.add(pipeline._kv_cache_memory_id("ref_image"))
        pipeline.transformer.evict_kv_cache_chunks(keep)

    @staticmethod
    def _to_uint8(decoded: torch.Tensor) -> np.ndarray:
        frames = (
            (decoded / 2 + 0.5).clamp(0, 1).mul_(255.0).round_().clamp_(0, 255)
            .to(torch.uint8)[0].permute(1, 2, 3, 0).contiguous()
        )
        return frames.cpu().numpy()

    # ------------------------------------------------------------ public loop
    @torch.no_grad()
    def process_chunk(self, source_frames: list[Image.Image]) -> np.ndarray:
        ref_chunk_latent = self._encode_reference_chunk(source_frames)
        latents = self._denoise_chunk(ref_chunk_latent)
        decoded = self._decode_chunk_pixels(latents)
        self._encode_pseudo_latent(decoded)
        self._evict_after_store()
        self.chunk_idx += 1
        return self._to_uint8(decoded)

    @torch.no_grad()
    def run(
        self,
        prompt: str,
        frames: list[Image.Image],
        ref_image: Image.Image | None = None,
    ) -> Iterator[tuple[int, np.ndarray]]:
        """Yield ``(chunk_idx, uint8 frames [T, H, W, 3])`` for every chunk."""
        if not frames:
            raise ValueError("No input frames.")
        self.start(prompt, frames[0], ref_image=ref_image)

        cursor = 0
        while cursor < len(frames):
            need = self.frames_per_next_chunk
            if len(frames) - cursor < need:
                break
            chunk_frames = frames[cursor: cursor + need]
            cursor += need
            idx = self.chunk_idx
            yield idx, self.process_chunk(chunk_frames)

    def close(self) -> None:
        self._clear_vae_caches()
        try:
            self.pipeline.transformer.reset_inference_kv_cache()
        except Exception:  # noqa: BLE001
            pass
        self.pseudo_latent = None
        self.prev_source_frame = None
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


__all__ = ["StreamingEditor", "StreamingSettings"]