import os
import subprocess
import shutil
import pickle
import cv2
import re

# 设置提示词并获取坐标
def set_prompt_and_get_coordinates(output_video, texts=['men', 'the table']):
    if isinstance(texts, str):
        texts = texts.split(',')
        texts = [text.strip() for text in texts]
        print(texts)
    with open('/kaggle/texts.pkl', 'wb') as file:
        pickle.dump(texts, file)
    with open('/kaggle/output_video2.pkl', 'wb') as file:
        pickle.dump(output_video, file)
    command = ['python', '/kaggle/florence-sam-kaggle/kaggle_florence_gpu_2.py']
    all_ok_bboxes = subprocess.run(command, capture_output=True, text=True)
    return all_ok_bboxes

# 运行 sam2 处理
def run_sam2(output_video):
    script_path = '/kaggle/florence-sam-kaggle/kaggle_sam2_gpu_2.py'
    command = ['python3', script_path]
    sam2_output = subprocess.run(command, capture_output=True, text=True)
    print(sam2_output)
    return sam2_output

# 生成带有音频的视频
def create_video_with_audio(image_folder, input_video_path):
    image_files = [f for f in os.listdir(image_folder) if f.endswith(('.png', '.jpg', '.jpeg'))]

    def natural_sort_key(s, _nsre=re.compile('([0-9]+)')):
        return [int(text) if text.isdigit() else text.lower() for text in re.split(_nsre, s)]

    image_files.sort(key=natural_sort_key)

    if image_files:
        first_image = cv2.imread(os.path.join(image_folder, image_files[0]))
        height, width, layers = first_image.shape
    else:
        raise ValueError("No valid images found in the folder after skipping the first one.")

    cap = cv2.VideoCapture(input_video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    cap.release()

    if fps <= 0:
        fps = 24  # 默认帧率

    output_video_path = os.path.join('/kaggle/working/sam2_videos/', os.path.basename(input_video_path))
    os.makedirs(os.path.dirname(output_video_path), exist_ok=True)

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    video_writer = cv2.VideoWriter('/kaggle/image_sequence_video2.mp4', fourcc, fps, (width, height))

    for image_file in image_files:
        image_path = os.path.join(image_folder, image_file)
        frame = cv2.imread(image_path)
        video_writer.write(frame)

    video_writer.release()

    temp_video_path = '/kaggle/image_sequence_video2.mp4'
    command = [
        'ffmpeg',
        '-y',  # 覆盖输出文件
        '-i', temp_video_path,
        '-i', input_video_path,
        '-c:v', 'copy',
        '-c:a', 'copy',
        '-shortest',
        output_video_path
    ]

    try:
        subprocess.run(command, check=True)
    except subprocess.CalledProcessError as e:
        print(f"Error running ffmpeg: {e}")
        print(f"Command: {' '.join(command)}")
        raise

    print(f"Video created successfully: {output_video_path}")
    return output_video_path

import os
import re

def natural_sort_key(s):
    """生成一个用于自然排序的键"""
    return [int(text) if text.isdigit() else text.lower() for text in re.split('([0-9]+)', s)]


# 处理所有视频文件
def process_all_videos(source_dir, target_dir, image_folder):
    os.makedirs(target_dir, exist_ok=True)
    video_files = [f for f in os.listdir(source_dir) if f.endswith(('.mp4', '.avi', '.mov'))]
    video_files.sort(key=natural_sort_key)
    # 反转列表
    video_files.reverse()
    print(video_files)
    for video_file in video_files:
        video_path = os.path.join(source_dir, video_file)
        print(f"Processing video: {video_path}")

        # 设置提示词并获取坐标
        result = set_prompt_and_get_coordinates(video_path, texts="men, the table")
        print(result)

        # 运行 sam2 处理
        result = run_sam2(video_path)
        print(result.stdout)

        # 生成带有音频的视频
        output_video_path = create_video_with_audio(image_folder, video_path)
        print(f"Output video: {output_video_path}")

        # 移动处理后的视频到目标目录
        target_video_path = os.path.join(target_dir, os.path.basename(output_video_path))
        shutil.move(video_path, target_video_path)
        print(f"Moved processed video to: {target_video_path}")

# 示例调用
source_dir = '/kaggle/o_videos'
target_dir = '/kaggle'
image_folder = '/kaggle/output2'

process_all_videos(source_dir, target_dir, image_folder)