# SPDX-License-Identifier: Apache-2.0
"""Sampling epoch for RL training — generate videos and
compute rewards."""

from __future__ import annotations

import time
from collections.abc import Callable
from typing import Any

import torch

from fastvideo.logger import init_logger
from fastvideo.train.methods.rl.utils.embeddings import (
    compute_text_embeddings,
)
from fastvideo.train.methods.rl.utils.pipeline import (
    wan_denoising_with_logprob,
)

logger = init_logger(__name__)

SEED_EPOCH_STRIDE = 10_000


def create_generator(
    prompts: list[str],
    base_seed: int,
    device: torch.device,
) -> list[torch.Generator]:
    """Create deterministic generators seeded by prompt."""
    generators = []
    for prompt in prompts:
        g = torch.Generator(device=device)
        g.manual_seed(base_seed + hash(prompt) % (2**31))
        generators.append(g)
    return generators


def sample_epoch(
    model,
    scheduler,
    train_sampler,
    train_iter,
    reward_fn: Callable,
    sample_neg_prompt_embeds: torch.Tensor,
    text_encoder,
    tokenizer,
    executor,
    epoch: int,
    global_step: int,
    *,
    # Config values passed explicitly.
    sample_batch_size: int,
    num_batches_per_epoch: int,
    num_inference_steps: int,
    guidance_scale: float,
    height: int,
    width: int,
    num_frames: int,
    noise_level: float,
    sde_type: str,
    diffusion_clip: bool,
    diffusion_clip_value: float,
    sde_window_size: int,
    sde_window_range: tuple[int, int] | None,
    kl_reward: float,
    same_latent: bool,
    seed: int,
    device: torch.device,
    is_main_process: bool,
    ref_transformer: torch.nn.Module | None = None,
    lora_model: Any | None = None,
    tracker: Any | None = None,
) -> tuple[
    list[dict[str, Any]],
    list[torch.Tensor],
    list[list[str]],
]:
    """Run one sampling epoch: generate videos, compute
    rewards asynchronously.

    Returns:
        Tuple of (samples, all_videos, all_prompts):
        - samples: list of sample dicts with prompt_ids,
          prompt_embeds, latents, log_probs, kl,
          timesteps, rewards.
        - all_videos: list of decoded video tensors
          per batch, each (B, 3, T, H, W) in [0, 1].
        - all_prompts: list of prompt string lists
          per batch.
    """
    samples = []
    all_videos: list[torch.Tensor] = []
    all_prompts: list[list[str]] = []

    for i in range(num_batches_per_epoch):
        current_epoch_tag = (
            epoch * num_batches_per_epoch + i
        )
        train_sampler.set_epoch(current_epoch_tag)

        # Drain until epoch tag matches.
        while True:
            epoch_tag, prompts, prompt_metadata = next(
                train_iter
            )
            if epoch_tag == current_epoch_tag:
                break

        torch.cuda.synchronize()
        _t_batch_start = time.perf_counter()

        _t_embed = time.perf_counter()
        prompt_embeds = compute_text_embeddings(
            prompts,
            text_encoder,
            tokenizer,
            max_sequence_length=512,
            device=device,
        )
        prompt_ids = tokenizer(
            prompts,
            padding="max_length",
            max_length=512,
            truncation=True,
            return_tensors="pt",
        ).input_ids.to(device)

        # Generator setup.
        if same_latent:
            gen = create_generator(
                prompts,
                base_seed=epoch * SEED_EPOCH_STRIDE + i,
                device=device,
            )
        else:
            gen = torch.Generator(device=device)
            gen.manual_seed(
                seed + epoch * SEED_EPOCH_STRIDE + i
            )

        torch.cuda.synchronize()
        _t_embed_done = time.perf_counter()

        with torch.no_grad():
            (
                videos,
                latents_list,
                log_probs_list,
                kls_list,
                timesteps_list,
            ) = wan_denoising_with_logprob(
                model,
                scheduler,
                prompt_embeds=prompt_embeds,
                negative_prompt_embeds=(
                    sample_neg_prompt_embeds
                ),
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                height=height,
                width=width,
                num_frames=num_frames,
                generator=gen,
                noise_level=noise_level,
                sde_type=sde_type,
                diffusion_clip=diffusion_clip,
                diffusion_clip_value=diffusion_clip_value,
                sde_window_size=sde_window_size,
                sde_window_range=sde_window_range,
                kl_reward=kl_reward,
                ref_transformer=ref_transformer,
                lora_model=lora_model,
            )

        torch.cuda.synchronize()
        _t_denoise_done = time.perf_counter()

        latents = torch.stack(latents_list, dim=1)
        log_probs = torch.stack(log_probs_list, dim=1)
        kls = torch.stack(kls_list, dim=1)
        kl = kls.detach()

        timesteps = (
            torch.stack(timesteps_list)
            .unsqueeze(0)
            .repeat(sample_batch_size, 1)
        )

        # Collect decoded videos and prompts for logging.
        all_videos.append(videos)
        all_prompts.append(list(prompts))

        # Async reward computation.
        rewards_future = executor.submit(
            reward_fn,
            videos,
            prompts,
            prompt_metadata,
            True,
        )
        time.sleep(0)

        logger.info(
            "[sample_epoch] batch %d/%d: "
            "text_embed=%.1fs  denoise=%.1fs  "
            "batch_total=%.1fs",
            i + 1,
            num_batches_per_epoch,
            _t_embed_done - _t_embed,
            _t_denoise_done - _t_embed_done,
            _t_denoise_done - _t_batch_start,
        )

        samples.append(
            {
                "prompt_ids": prompt_ids,
                "prompt_embeds": prompt_embeds,
                "negative_prompt_embeds": (
                    sample_neg_prompt_embeds
                ),
                "timesteps": timesteps,
                "latents": latents[:, :-1],
                "next_latents": latents[:, 1:],
                "log_probs": log_probs,
                "kl": kl,
                "rewards": rewards_future,
            }
        )

    # Wait for all rewards.
    torch.cuda.synchronize()
    _t_reward_wait = time.perf_counter()
    for sample in samples:
        rewards, _ = sample["rewards"].result()
        sample["rewards"] = {
            key: torch.as_tensor(value, device=device).float()
            for key, value in rewards.items()
        }
    _t_reward_done = time.perf_counter()
    logger.info(
        "[sample_epoch] reward_wait=%.1fs",
        _t_reward_done - _t_reward_wait,
    )

    return samples, all_videos, all_prompts