# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import os
import argparse
from process_pipepline import ProcessPipeline

# simple args container like before
class _Args:
    pass

def _parse_args():
    args = _Args()

    # general paths
    args.ckpt_path     = "./Wan2.2-Animate-14B/process_checkpoint"
    args.video_path    = None
    args.refer_path    = None
    args.save_path     = None

    # processing parameters
    args.resolution_area = [1280, 720]
    # Keep in sync with wan_engine.FPS. This defaulted to 16 while wan_engine
    # defaulted to 8 and NOTHING set the env var, so get_frames built 160 frame
    # indices over 80 unique frames — every frame processed twice, then exported
    # at 8fps as a 20s half-speed video. wan_engine now force-sets ANIMATE_FPS at
    # import so the two can never diverge again; 8 is the matching default.
    args.fps             = int(os.getenv("ANIMATE_FPS", "8"))

    # feature flags
    args.replace_flag  = True
    args.retarget_flag = False
    args.use_flux      = False

    # mask strategy parameters (replacement mode)
    args.iterations = 3
    args.k          = 7
    args.w_len      = 1
    args.h_len      = 1

    return args

def load_preprocess_models(with_pose=True, pose_device="cuda"):
    """with_pose=False builds a SAM2-only pipeline. Pose now runs on the CPU in
    the main process (wan_engine._pose_cpu), so rebuilding the 2.5GB ViTPose ONNX
    session inside the paid GPU reservation just to leave it idle is pure waste."""
    ckpt_path = "./Wan2.2-Animate-14B/process_checkpoint"

    pose2d_checkpoint_path = (
        os.path.join(ckpt_path, 'pose2d/vitpose_h_wholebody.onnx') if with_pose else None)
    det_checkpoint_path = (
        os.path.join(ckpt_path, 'det/yolov10m.onnx') if with_pose else None)

    sam2_checkpoint_path = [os.path.join(ckpt_path, 'sam2/sam2_hiera_large.pt'),"sam2_hiera_l.yaml"]

    flux_kontext_path = None

    process_pipeline = ProcessPipeline(det_checkpoint_path=det_checkpoint_path, pose2d_checkpoint_path=pose2d_checkpoint_path, sam_checkpoint_path=sam2_checkpoint_path, flux_kontext_path=flux_kontext_path, pose_device=pose_device)

    return process_pipeline

def run(process_pipeline, input_video, edited_frame, preprocess_dir, w, h, tag_string,
        pts_by_frame: dict, lbs_by_frame: dict,
        frames=None, precomputed_metas=None):
    args = _parse_args()

    if tag_string == "retarget_flag":
        retarget_flag = True
        replace_flag = False
    else:
        retarget_flag = False
        replace_flag = True

    src_pose_path, src_face_path, src_bg_path, src_mask_path, src_ref_image = process_pipeline(video_path=input_video, 
                     refer_image_path=edited_frame, 
                     output_path=preprocess_dir,
                     resolution_area=[w, h],
                     fps=args.fps,
                     iterations=args.iterations,
                     k=args.k,
                     w_len=args.w_len,
                     h_len=args.h_len,
                     retarget_flag=retarget_flag,
                     use_flux=args.use_flux,
                     replace_flag=replace_flag,
                     pts_by_frame=pts_by_frame,
                     lbs_by_frame=lbs_by_frame,
                     frames=frames,
                     precomputed_metas=precomputed_metas)
    
    return src_pose_path, src_face_path, src_bg_path, src_mask_path, src_ref_image