# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import argparse
import logging
import os
import sys
import warnings
from datetime import datetime
import cv2
warnings.filterwarnings('ignore')
from wan.timeutils import ClockContext
import random
import torch
import torch.distributed as dist
from PIL import Image

import wan
from wan.configs import MAX_AREA_CONFIGS, SIZE_CONFIGS, SUPPORTED_SIZES, WAN_CONFIGS
from wan.distributed.util import init_distributed_group
from wan.utils.utils import save_video, str2bool
from wan.bench_speed.textimage2video import WanTI2V
from generate_causal import _validate_args, _parse_args, _init_logging

EXAMPLE_PROMPT = {
    "ti2v-5B": {
        "prompt":
            "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.",
    },
}


def generate(args, new_args):
    assert not args.use_prompt_extend, "Prompt extension is not supported in WanTI2V generation."
    rank = 0
    world_size = 1
    local_rank = 0
    device = local_rank
    _init_logging(rank)

    if args.offload_model is None:
        args.offload_model = False if world_size > 1 else True
        logging.info(
            f"offload_model is not specified, set to {args.offload_model}.")
    if world_size > 1:
        torch.cuda.set_device(local_rank)
        dist.init_process_group(
            backend="nccl",
            init_method="env://",
            rank=rank,
            world_size=world_size)
    else:
        assert not (
            args.t5_fsdp or args.dit_fsdp
        ), f"t5_fsdp and dit_fsdp are not supported in non-distributed environments."
        assert not (
            args.ulysses_size > 1
        ), f"sequence parallel are not supported in non-distributed environments."

    if args.ulysses_size > 1:
        assert args.ulysses_size == world_size, f"The number of ulysses_size should be equal to the world size."
        init_distributed_group()

    cfg = WAN_CONFIGS[args.task]
    if args.ulysses_size > 1:
        assert cfg.num_heads % args.ulysses_size == 0, f"`{cfg.num_heads=}` cannot be divided evenly by `{args.ulysses_size=}`."

    logging.info(f"Generation job args: {args}")
    logging.info(f"Generation model config: {cfg}")

    if dist.is_initialized():
        base_seed = [args.base_seed] if rank == 0 else [None]
        dist.broadcast_object_list(base_seed, src=0)
        args.base_seed = base_seed[0]

    logging.info(f"Input prompt: {args.prompt}")


    logging.info("Creating WanTI2V pipeline.")
    wan_ti2v = WanTI2V(
        config=cfg,
        checkpoint_dir=args.ckpt_dir,
        pt_dir=args.pt_dir,
        device_id=device,
        rank=rank,
        t5_fsdp=args.t5_fsdp,
        dit_fsdp=args.dit_fsdp,
        use_sp=(args.ulysses_size > 1),
        t5_cpu=args.t5_cpu,
        convert_model_dtype=args.convert_model_dtype,
    )

    logging.info(f"Generating video ...")
    txt_files = os.listdir(os.path.join(new_args.dataset_path, "metas"))
    random.seed(42)
    random.shuffle(txt_files)
    for txt_file in txt_files[:20]:
        output_prefix = txt_file.replace(".txt", "")
        mp4_path = os.path.join(new_args.dataset_path, 'videos', txt_file.replace(".txt", ".mp4"))
        with open(os.path.join(new_args.dataset_path, "metas", txt_file), "r") as f:
            prompt = f.read().strip()
            print(prompt)
        if mp4_path.endswith('.mp4'):
            cap = cv2.VideoCapture(mp4_path)
            ret, frame = cap.read()
            if ret:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                img = Image.fromarray(frame)
            cap.release()
        else:
            img = Image.open(mp4_path)
        img = img.resize(SIZE_CONFIGS[args.size])
        with ClockContext(f"{'end to end gen':-^30}"):
            video = wan_ti2v.generate(
                prompt,
                img=img,
                size=SIZE_CONFIGS[args.size],
                max_area=MAX_AREA_CONFIGS[args.size],
                frame_num=args.frame_num,
                shift=args.sample_shift,
                sample_solver=args.sample_solver,
                sampling_steps=args.sample_steps,
                guide_scale=args.sample_guide_scale,
                seed=args.base_seed,
                offload_model=args.offload_model
            )


        if rank == 0:
            formatted_time = datetime.now().strftime("%Y%m%d_%H%M%S")
            formatted_prompt = output_prefix.replace(" ", "_").replace("/","_")[:70]
            suffix = '.mp4'
            # save_file_path = f"{args.task}_{args.size.replace('*','x') if sys.platform=='win32' else args.size}_{args.ulysses_size}_{formatted_prompt}_{formatted_time}" + suffix
            save_file_path = f"{args.task}_{args.size.replace('*','x') if sys.platform=='win32' else args.size}_{args.ulysses_size}_{formatted_prompt}" + suffix
            save_file_path = os.path.join(new_args.output_dir, save_file_path)

            logging.info(f"Saving generated video to {save_file_path}")
            save_video(
                tensor=video[None],
                save_file=save_file_path,
                fps=10,
                nrow=1,
                normalize=True,
                value_range=(-1, 1))
    del video

    torch.cuda.synchronize()
    if dist.is_initialized():
        dist.barrier()
        dist.destroy_process_group()

    logging.info("Finished.")


if __name__ == "__main__":
    args, left_args = _parse_args()
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset_path",type=str, default=None)
    parser.add_argument("--output_dir",type=str, required=True)
    new_args = parser.parse_args(left_args)
    os.makedirs(new_args.output_dir, exist_ok=True)
    generate(args, new_args)
