import os
import cv2
import numpy as np
from PIL import Image
import argparse

def is_image_file(filename):
    return filename.lower().endswith(('.jpg', '.jpeg', '.png'))

def resize_and_pad_cv(image, size=(1280, 720)):
    ih, iw = image.shape[:2]
    target_w, target_h = size

    scale = min(target_w / iw, target_h / ih)
    new_w = int(iw * scale)
    new_h = int(ih * scale)
    resized = cv2.resize(image, (new_w, new_h))

    top = (target_h - new_h) // 2
    bottom = target_h - new_h - top
    left = (target_w - new_w) // 2
    right = target_w - new_w - left

    padded = cv2.copyMakeBorder(resized, top, bottom, left, right, cv2.BORDER_CONSTANT, value=[0, 0, 0])
    return padded

def get_first_frame(video_path):
    cap = cv2.VideoCapture(video_path)
    success, frame = cap.read()
    cap.release()
    return frame if success else None

def check_and_process_individual_cover(mp4_path):
    base = os.path.splitext(mp4_path)[0]
    for ext in ['.jpg', '.jpeg', '.png']:
        img_path = base + ext
        if os.path.exists(img_path):
            try:
                with Image.open(img_path) as img:
                    w, h = img.size
                    if w < 720:
                        os.remove(img_path)
                        print(f"删除低分辨率封面图: {img_path}")
                        return False  # 继续处理 mp4
                    else:
                        print(f"已有封面图，裁剪为1280x720并跳过视频: {img_path}")
                        processed = resize_and_pad_cv(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))
                        #cv2.imwrite(img_path, processed)
                        Image.fromarray(cv2.cvtColor(processed, cv2.COLOR_BGR2RGB)).save(img_path)

                        return True  # 不再处理该 mp4
            except Exception as e:
                print(f"处理封面图出错 {img_path}: {e}")
                return False
    return False  # 没有封面图，继续处理 mp4

def process_medias(folder):
    for root, _, files in os.walk(folder):
        for file in files:
            if file.lower().endswith('.mp4'):
                video_path = os.path.join(root, file)
                skip = check_and_process_individual_cover(video_path)
                if skip:
                    continue
                frame = get_first_frame(video_path)
                if frame is None:
                    print(f"无法读取视频帧: {video_path}")
                    continue
                output_img = resize_and_pad_cv(frame)
                output_path = os.path.splitext(video_path)[0] + "_cover.jpg"
                #cv2.imwrite(output_path, output_img)
                Image.fromarray(cv2.cvtColor(output_img, cv2.COLOR_BGR2RGB)).save(output_path)
                print(f"生成封面图: {output_path}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="处理mp4封面图")
    parser.add_argument("folder", help="包含视频的文件夹")
    args = parser.parse_args()
    process_medias(args.folder)