# 下載VideoSubFinder, 改path_VideoSubFinder
# 裝CUDA 12.9
# pip install paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu129/
# pip install paddleocr

import os
import glob
import subprocess
import time
import argparse
import numpy as np
from PIL import Image
from paddleocr import PaddleOCR

path_VideoSubFinder = r"D:\work\workspace\SubtitleOCR\VideoSubFinder"
fn_VideoSubFinder = "VideoSubFinderWXW_intel.exe"
path_VideoSubFinder_exe = fr"{path_VideoSubFinder}\{fn_VideoSubFinder}"
path_sub_images = fr"{path_VideoSubFinder}\RGBImages"

parser = argparse.ArgumentParser()
parser.add_argument('-i', '--image', required=False, help='image file path')
parser.add_argument('-v', '--video', required=False, help='video file path')
parser.add_argument('-te', '--sub_top', type=float, default=0.15, help='top of subtitle, percent offset of height, from bottom. (default=0.15)')
parser.add_argument('-be', '--sub_bottom', type=float, default=0.05, help='bottom of subtitle, percent offset of height, from bottom. (default=0.05)')
parser.add_argument('-mw', '--ocr_min_width', type=int, default=100, help='ocr: minimum text width. (default=100)')
parser.add_argument('-mh', '--ocr_min_height', type=int, default=50, help='ocr: minimum text height. (default=50)')
args = parser.parse_args()

print("Initializing PaddleOCR...")
ocr = PaddleOCR(
    text_detection_model_name="PP-OCRv5_server_det",
    text_recognition_model_name="PP-OCRv5_server_rec",
    use_doc_orientation_classify=False,
    use_doc_unwarping=False,
    use_textline_orientation=False,
)

def do_ocr(img_path):
    try:
        img = np.array(Image.open(img_path).convert('RGB'))
    except Exception as e:
        print(f"[ERROR] Failed to open image {img_path}: {e}")
        return ""

    result = ocr.predict(img)
    for res in result:
        if 'rec_texts' not in res or 'rec_polys' not in res:
            continue
        txts, polys = res['rec_texts'], res['rec_polys']
        kept_txts = []
        for txt, poly in zip(txts, polys):
            xs = [p[0] for p in poly]
            ys = [p[1] for p in poly]
            width, height = max(xs) - min(xs), max(ys) - min(ys)
            if height >= args.ocr_min_height and width >= args.ocr_min_width:
                kept_txts.append(txt)
        if kept_txts:
            return " ".join(kept_txts)
    return ""

# If only an image is provided, do single OCR and exit
if args.image and not args.video:
    print(f"Performing OCR on image: {args.image}")
    txt = do_ocr(args.image)
    print(txt)
    print("Single image OCR complete.")
    exit(0)

if not args.video:
    print("[ERROR] Missing required argument: -v/--video when not using -i/--image only.")
    exit(1)

# Step 1: Run VideoSubFinder
print("Step 1: Running VideoSubFinder...")
params = [path_VideoSubFinder_exe, '-uc', '-c', '-r', '-i', args.image or args.video, '-te', str(args.sub_top), '-be', str(args.sub_bottom)]
proc = subprocess.Popen(params)

print("Waiting for VideoSubFinder to complete...")
flag_first = True
path_sub_images_jpeg = os.path.join(path_sub_images, '*.jpeg')
while proc.poll() is None:
    time.sleep(3)
    jpegs = glob.glob(path_sub_images_jpeg)
    if jpegs:
        print(jpegs[-1], flush=True)
        if flag_first:
            os.startfile(jpegs[0])
            flag_first = False
print("VideoSubFinder finished.")


# Step 2: Get subtitle image files
print("Step 2: Collecting subtitle images...")
imgs = sorted(glob.glob(os.path.join(path_sub_images, '*.jpeg')))
print(f"Found {len(imgs)} subtitle images.")

# Step 3: OCR each subtitle image and generate VTT
print("Step 3: Performing OCR and writing .vtt file...")
output_vtt = f"{args.video}.vtt"
with open(output_vtt, 'w', encoding='utf-8') as f:
    f.write("WEBVTT\n\n")
    for img in imgs:
        base = os.path.basename(img)
        parts = base.split('__')
        if len(parts) < 2:
            continue
        start, end = parts[0], parts[1].split('_000000')[0]
        start = start.replace('_', ':', 2).replace('_', '.')
        end = end.replace('_', ':', 2).replace('_', '.')
        txt = do_ocr(img)
        if not txt:
            continue
        print(f"[{time.strftime('%H%M%S')}] {start} --> {end}: {txt}")
        f.write(f"{start} --> {end}\n{txt}\n\n")

print(f"Subtitle extraction complete: {output_vtt}")