import os
import subprocess
import pickle
import cv2
from moviepy.editor import VideoFileClip, ImageSequenceClip
import shutil
import re
# 设置环境变量
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

def uninstall_tensorflow():
    try:
        result = subprocess.run(['pip', 'uninstall', '-y', 'tensorflow'], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        print("卸载成功")
        print("输出:", result.stdout.decode('utf-8'))
    except subprocess.CalledProcessError as e:
        print("卸载失败")
        print("错误:", e.stderr.decode('utf-8'))

# 调用函数
uninstall_tensorflow()

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('/workspace/florence-sam/texts.pkl', 'wb') as file:
        pickle.dump(texts, file)
    with open('/workspace/florence-sam/output_video.pkl', 'wb') as file:
        pickle.dump(output_video, file)
    command = ['python', '/workspace/florence-sam/1.py']
    all_ok_bboxes = subprocess.run(command, capture_output=True, text=True)
    return all_ok_bboxes

def run_sam2(output_video):
    script_path = '/workspace/florence-sam/2.py'
    command = ['python3', script_path]
    sam2_output = subprocess.run(command, capture_output=True, text=True)
    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()

    image_paths = [os.path.join(image_folder, img) for img in image_files]
    clip = ImageSequenceClip(image_paths, fps=fps)

    audio_clip = VideoFileClip(input_video_path).audio
    final_clip = clip.set_audio(audio_clip)

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

    final_clip.write_videofile(output_video_path, codec='libx264')

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

def process_videos(source_dir, target_dir, texts="men, the table"):
    os.makedirs(target_dir, exist_ok=True)

    for video_file in os.listdir(source_dir):
        if video_file.endswith(('.mp4', '.avi', '.mov')):
            video_path = os.path.join(source_dir, video_file)
            print(f"Processing video: {video_path}")

            result = set_prompt_and_get_coordinates(video_path, texts)
            print(result)

            result = run_sam2(video_path)
            print(result)

            with open('/workspace/florence-sam/output_video.pkl', 'rb') as file:
                output_video = pickle.load(file)

            create_video_with_audio('/workspace/output', output_video)

            # 移动原视频到目标目录
            shutil.move(video_path, os.path.join(target_dir, video_file))

if __name__ == "__main__":
    source_dir = '/workspace/transvnet/o_videos'
    target_dir = '/workspace/alreayvideos'
    texts = "men, the table"
    process_videos(source_dir, target_dir, texts)