from __future__ import annotations

import os
import time
from pathlib import Path

import numpy as np
import torch

from common import MODEL_REPO, drop_cached, load_wire, snapshot

NAME = "diffusers"
LABEL = "diffusers (PR #2)"

STATE = {}


def load(partition: str) -> None:
    from diffusers.modular_pipelines.pika_h3_flex import PikaH3FlexBlocks

    trunk = "transformer_ref" if partition == "ref2va" else "transformer"
    started = time.time()
    blocks = PikaH3FlexBlocks()
    blocks.sub_blocks.pop("text_encoder")
    pipe = blocks.init_pipeline(MODEL_REPO)
    pipe.load_components(names=["vae", "audio_vae"], dtype=torch.float32)
    pipe.load_components(names=[trunk, "scheduler", "audio_scheduler"], dtype=torch.bfloat16)
    adapter = Path(snapshot([f"{partition}/*"])) / partition / "adapter.safetensors"
    pipe.load_lora_weights(
        str(adapter),
        **({"load_into_transformer_ref": True} if partition == "ref2va" else {}),
    )
    drop_cached(adapter)
    getattr(pipe, trunk).requires_grad_(False).eval()
    attention = os.environ.get("PIKA_ATTENTION")
    if attention:
        getattr(pipe, trunk).set_attention_backend(attention)
    getattr(pipe, trunk).to("cuda")
    STATE.update(partition=partition, trunk=trunk, pipe=pipe, loaded_in=time.time() - started)
    print(f"[diffusers] {partition}: {trunk} + {partition} adapter ready in {STATE['loaded_in']:.0f}s", flush=True)


def video_vae():
    return STATE["pipe"].vae


def transformer():
    return getattr(STATE["pipe"], STATE["trunk"])


def prepare(task, prompt, first_image, last_image, references, height, width, num_frames):
    from diffusers.modular_pipelines.pika_h3_flex import (
        PikaH3FlexAudioReference,
        PikaH3FlexImageReference,
        PikaH3FlexVideoReference,
    )
    from PIL import Image

    kinds = {"image": PikaH3FlexImageReference, "video": PikaH3FlexVideoReference, "audio": PikaH3FlexAudioReference}
    inputs = {"height": height, "width": width, "num_frames": num_frames}
    if task in ("i2va", "fl2va") and first_image:
        inputs["image"] = Image.open(first_image).convert("RGB")
    if task in ("l2va", "fl2va") and last_image:
        inputs["last_image"] = Image.open(last_image).convert("RGB")
    if task == "ref2va":
        inputs["references"] = [kinds[kind].from_file(path) for kind, path in references]
    return inputs


def run(prepared, wire_path, seed, progress=None, noise=None):
    pipe = STATE["pipe"]
    prompt_embeds, text_token_tags, _ = load_wire(wire_path)
    inputs = dict(prepared)
    if noise is not None:
        inputs["latents"], inputs["audio_latents"] = noise
    pipe.vae.to("cuda")
    pipe.audio_vae.to("cuda")
    if progress is not None:
        calls = {"n": 0}

        def hook(module, args, kwargs):
            calls["n"] += 1
            progress(calls["n"] / 6, desc=f"Flex step {calls['n']} of 6")

        handle = transformer().register_forward_pre_hook(hook, with_kwargs=True)
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    start = time.perf_counter()
    try:
        with torch.inference_mode():
            state = pipe(
                prompt_embeds=prompt_embeds.to("cuda"),
                text_token_tags=text_token_tags,
                generator=torch.Generator("cuda").manual_seed(int(seed)),
                output_type="np",
                **inputs,
            )
        torch.cuda.synchronize()
    finally:
        if progress is not None:
            handle.remove()
    total = time.perf_counter() - start
    videos = state.get("videos")
    frames = np.asarray(videos[0])
    if frames.dtype != np.uint8:
        frames = (np.clip(frames, 0.0, 1.0) * 255.0).round().astype(np.uint8)
    audio = state.get("audio")[0]
    wave = audio.float().cpu().numpy() if torch.is_tensor(audio) else np.asarray(audio, dtype=np.float32)
    timings = {"pipeline": total, "peak_gib": torch.cuda.max_memory_allocated() / 2**30}
    return frames, wave, int(state.get("sampling_rate")), timings, None