import sys
import os
import cv2
import time
import difflib
import argparse
import traceback
import json
from pathlib import Path

def register_ocr_gpu_dll_paths():
    import os
    from pathlib import Path

    root = Path(__file__).resolve().parents[2]
    candidates = [
        root / "env" / "Lib" / "site-packages" / "nvidia" / "cudnn" / "bin",
        root / "env" / "Lib" / "site-packages" / "nvidia" / "cublas" / "bin",
        root / "env" / "Lib" / "site-packages" / "nvidia" / "cuda_runtime" / "bin",
        root / "env" / "Lib" / "site-packages" / "nvidia" / "cuda_nvrtc" / "bin",
        root / "env" / "Lib" / "site-packages" / "paddle" / "libs",
    ]

    added = []
    for p in candidates:
        if p.exists():
            try:
                os.add_dll_directory(str(p))
            except Exception as e:
                print(f"[OCR DLL WARN] add_dll_directory failed for {p}: {e}")
            added.append(str(p))

    if added:
        os.environ["PATH"] = os.pathsep.join(added) + os.pathsep + os.environ.get("PATH", "")

    print("[OCR DLL] registered paths:")
    for p in added:
        print(f"  - {p}")

register_ocr_gpu_dll_paths()

# Enforce UTF-8 for Windows console
if sys.platform == 'win32':
    try:
        if hasattr(sys.stdout, 'reconfigure'):
            sys.stdout.reconfigure(encoding='utf-8')
        if hasattr(sys.stderr, 'reconfigure'):
            sys.stderr.reconfigure(encoding='utf-8')
    except Exception:
        pass

def get_similarity(s1, s2):
    if not s1 or not s2:
        return 0.0
    return difflib.SequenceMatcher(None, s1, s2).ratio()

def _is_cjk_char(ch):
    return "\u3400" <= ch <= "\u9fff"

def _ocr_garbage_reason(text, max_frame_chars=120):
    text = " ".join(str(text or "").split()).strip()
    if not text:
        return "empty"
    if len(text) > int(max_frame_chars):
        return "frame_text_too_long"

    chars = [c for c in text if not c.isspace()]
    if not chars:
        return "empty"

    noise_chars = set("□√×÷′²③④")
    noise_ratio = sum(1 for c in chars if c in noise_chars) / len(chars)
    if noise_ratio > 0.25:
        return "ocr_symbol_noise"

    cjk_chars = [c for c in chars if _is_cjk_char(c)]
    if len(cjk_chars) >= 40:
        counts = {}
        for c in cjk_chars:
            counts[c] = counts.get(c, 0) + 1
        dominant_ratio = max(counts.values()) / len(cjk_chars)
        unique_ratio = len(counts) / len(cjk_chars)
        if dominant_ratio > 0.35:
            return "dominant_cjk_repeat"
        if len(cjk_chars) >= 80 and unique_ratio < 0.15:
            return "low_cjk_diversity"

    return None

def _atomic_write_json(path, payload):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = path.with_name(path.name + ".tmp")
    with open(tmp_path, "w", encoding="utf-8") as f:
        json.dump(payload, f, ensure_ascii=False, indent=2)
    os.replace(tmp_path, path)

def _atomic_write_text(path, content):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = path.with_name(path.name + ".tmp")
    with open(tmp_path, "w", encoding="utf-8") as f:
        f.write(content)
    os.replace(tmp_path, path)

def _extract_ocr_lines(results, min_confidence=0.0):
    lines = []
    for page in (results or []):
        if not page:
            continue
        for line in page:
            if not line or len(line) <= 1 or not line[1]:
                continue
            try:
                text = str(line[1][0] or "").strip()
                score = float(line[1][1])
            except Exception:
                continue
            if text and score >= float(min_confidence):
                lines.append((text, score))
    return lines

def _find_cjk_font():
    candidates = [
        Path("C:/Windows/Fonts/msyh.ttc"),
        Path("C:/Windows/Fonts/msyhbd.ttc"),
        Path("C:/Windows/Fonts/simhei.ttf"),
        Path("C:/Windows/Fonts/simsun.ttc"),
    ]
    for path in candidates:
        if path.exists():
            return path
    return None

def _build_ocr_probe_image():
    from PIL import Image, ImageDraw, ImageFont
    import numpy as np

    text = "现在主要是拉尔夫劳伦贸易"
    image = Image.new("RGB", (760, 128), (22, 22, 22))
    draw = ImageDraw.Draw(image)
    font_path = _find_cjk_font()
    if not font_path:
        raise RuntimeError("No Windows CJK font found for OCR quality probe.")
    font = ImageFont.truetype(str(font_path), 46)
    draw.text((28, 32), text, font=font, fill=(245, 245, 245))
    return np.array(image), text

def _run_ocr_quality_probe(reader, max_frame_chars=160):
    probe_img, expected_text = _build_ocr_probe_image()
    results = reader.ocr(probe_img, cls=True)
    lines = _extract_ocr_lines(results, min_confidence=0.1)
    joined = " ".join(text for text, _score in lines).strip()
    joined = " ".join(joined.split())
    expected_chars = set("现在主要拉尔夫劳伦贸易")
    hit_count = sum(1 for ch in expected_chars if ch in joined)
    reason = _ocr_garbage_reason(joined, max_frame_chars=max_frame_chars)
    ok = bool(joined) and hit_count >= 4 and not reason and len(lines) <= 8
    details = {
        "ok": ok,
        "expected": expected_text,
        "recognized": joined[:240],
        "line_count": len(lines),
        "hit_count": hit_count,
        "reject_reason": reason,
        "sample_lines": [
            {"text": text, "confidence": round(float(score), 4)}
            for text, score in lines[:5]
        ],
    }
    return ok, details

def format_srt_time(seconds):
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = int(seconds % 60)
    ms = int((seconds - int(seconds)) * 1000)
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"

def is_cuda_fully_functional():
    try:
        import torch
        if not torch.cuda.is_available():
            return False
            
        device_idx = torch.cuda.current_device()
        gpu_major, gpu_minor = torch.cuda.get_device_capability(device_idx)
        arch_list = torch.cuda.get_arch_list()
        if not arch_list:
            return False
            
        supported_caps = []
        for arch in arch_list:
            digits = ''.join(c for c in arch if c.isdigit())
            if digits:
                val = int(digits)
                major = val // 10
                minor = val % 10
                supported_caps.append((major, minor))
                
        if not supported_caps:
            return False
            
        max_supported_major = max(major for major, minor in supported_caps)
        if gpu_major > max_supported_major:
            return False
            
        import torch.nn as nn
        conv = nn.Conv2d(1, 1, 3).cuda()
        x = torch.randn(1, 1, 8, 8).cuda()
        y = conv(x)
        float(y.to("cpu")[0, 0, 0, 0])
        return True
    except Exception:
        return False

def main():
    parser = argparse.ArgumentParser(description="Standalone PaddleOCR Subtitle Extractor CLI")
    parser.add_argument("--video", required=True, help="Path to input video file")
    parser.add_argument("--output", required=True, help="Path to output SRT file")
    parser.add_argument("--region", required=False, help="Selected subtitle region as 'x,y,w,h,orig_w,orig_h'")
    parser.add_argument("--device", default="auto", choices=["cuda", "cpu", "auto"], help="Device to run inference on")
    parser.add_argument("--prefer-gpu", default="true", help="Prefer GPU execution (true/false)")
    parser.add_argument("--allow-cpu-fallback", default="true", help="Allow fallback to CPU (true/false)")
    parser.add_argument("--gpu-mem-mb", type=int, default=3072, help="GPU memory allocation limit in MB")
    parser.add_argument("--sample-interval", type=float, default=0.25, help="Sample interval in seconds")
    parser.add_argument("--min-line-confidence", type=float, default=0.45, help="Drop OCR text lines below this confidence")
    parser.add_argument("--max-frame-text-chars", type=int, default=120, help="Reject OCR frame text longer than this")
    parser.add_argument("--max-block-duration", type=float, default=6.5, help="Maximum subtitle block duration before forced split")
    args = parser.parse_args()

    video_path = Path(args.video)
    output_srt_path = Path(args.output)
    
    if not video_path.exists():
        print(f"Error: Video file not found at {video_path}", file=sys.stderr)
        sys.exit(2)

    prefer_gpu = args.prefer_gpu.lower() in ("true", "1", "yes", "t")
    allow_cpu_fallback = args.allow_cpu_fallback.lower() in ("true", "1", "yes", "t")

    # ── Step 1: Set environment variables BEFORE any paddle/paddleocr import ──
    import os
    os.environ.setdefault("PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION", "python")
    os.environ["FLAGS_use_mkldnn"] = "0"
    print("[OCR ENV] PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python")
    try:
        import google.protobuf
        pb_ver = google.protobuf.__version__
    except Exception:
        pb_ver = "Unknown"
    print(f"[OCR ENV] protobuf version: {pb_ver}")

    # ── Step 2: Import paddle (lightweight) and check CUDA BEFORE loading PaddleOCR ──
    try:
        import paddle
    except ImportError as e:
        print(f"Error: paddle not installed. {e}", file=sys.stderr)
        sys.exit(1)

    compiled_with_cuda = False
    try:
        compiled_with_cuda = paddle.device.is_compiled_with_cuda()
    except Exception:
        pass

    # Log configuration and GPU check
    print(f"[OCR CONFIG] requested device: {args.device}")
    print(f"[OCR CONFIG] allow_cpu_fallback: {allow_cpu_fallback}")
    print(f"[OCR GPU CHECK] paddle compiled with cuda: {compiled_with_cuda}")
    try:
        paddle_device = paddle.device.get_device()
    except Exception:
        paddle_device = "unknown"
    print(f"[OCR GPU CHECK] paddle current device: {paddle_device}")

    # ── Step 3: FAIL-FAST gate — exit before any heavy model loading ──
    run_on_gpu = False
    fallback_reason = None

    if args.device == "cuda":
        if not compiled_with_cuda:
            fallback_reason = "Paddle is not compiled with CUDA."
            if not allow_cpu_fallback:
                print("[OCR ERROR] PaddleOCR GPU requested but Paddle is CPU-only. Failing early.", file=sys.stderr)
                print(f"Reason: {fallback_reason}", file=sys.stderr)
                sys.exit(3)
            run_on_gpu = False
        else:
            run_on_gpu = True
    elif args.device == "cpu":
        run_on_gpu = False
    else: # auto
        if compiled_with_cuda and prefer_gpu:
            run_on_gpu = True
        else:
            if not compiled_with_cuda:
                fallback_reason = "Paddle is not compiled with CUDA."
            else:
                fallback_reason = "GPU not preferred by configuration."
            if not allow_cpu_fallback:
                print("[OCR ERROR] PaddleOCR GPU requested but Paddle is CPU-only. Failing early.", file=sys.stderr)
                print(f"Reason: {fallback_reason}", file=sys.stderr)
                sys.exit(3)
            run_on_gpu = False

    # Log device selection
    if run_on_gpu:
        print("[OCR] selected device: gpu")
    else:
        print(f"[OCR WARNING] PaddleOCR GPU unavailable, falling back to CPU.")
        print(f"Reason: {fallback_reason or 'User requested CPU execution.'}")

    if run_on_gpu:
        try:
            paddle.set_device("gpu:0")
        except Exception as e:
            print(f"Warning: paddle.set_device('gpu:0') failed: {e}")

    # ── Step 4: NOW import PaddleOCR (heavy — downloads/loads models) ──
    print("Loading PaddleOCR model...")
    try:
        from paddleocr import PaddleOCR
    except ImportError as e:
        print(f"Error: paddleocr not installed. {e}", file=sys.stderr)
        sys.exit(1)

    # Inspect PaddleOCR constructor signature
    import inspect
    sig = inspect.signature(PaddleOCR)
    print(f"[DEBUG] PaddleOCR constructor signature: {sig}")

    ocr_params = {
        "use_angle_cls": True,
        "lang": 'ch',
        "show_log": False,
        "enable_mkldnn": False
    }

    if "device" in sig.parameters:
        ocr_params["device"] = "gpu" if run_on_gpu else "cpu"
    else:
        ocr_params["use_gpu"] = run_on_gpu
        if run_on_gpu:
            ocr_params["gpu_mem"] = args.gpu_mem_mb

    try:
        reader = PaddleOCR(**ocr_params)
    except Exception as e:
        if not allow_cpu_fallback and run_on_gpu:
            print("PaddleOCR GPU requested but unavailable.", file=sys.stderr)
            print(f"Reason: PaddleOCR GPU init failed (Error: {e}).", file=sys.stderr)
            sys.exit(3)
        print(f"Warning: Failed to load PaddleOCR on GPU. Falling back to CPU. Error: {e}")
        # Try CPU fallback
        if "device" in sig.parameters:
            ocr_params["device"] = "cpu"
        else:
            ocr_params["use_gpu"] = False
            if "gpu_mem" in ocr_params:
                del ocr_params["gpu_mem"]
        try:
            reader = PaddleOCR(**ocr_params)
        except Exception as ex:
            print(f"PaddleOCR initialization failed.\nReason: {ex}", file=sys.stderr)
            sys.exit(4)

    # ── OCR GPU Preflight Inference Check ──
    if run_on_gpu:
        print("[OCR GPU PREFLIGHT] running first inference test...")
        try:
            import numpy as np
            dummy_img = np.zeros((64, 64, 3), dtype=np.uint8)
            reader.ocr(dummy_img, cls=True)
            print("[OCR GPU PREFLIGHT] PASS")
        except Exception as e:
            err_msg = str(e)
            if "cudnn64_8.dll" in err_msg or "PreconditionNotMet" in err_msg or "dynamic library" in err_msg or "DLL load failed" in err_msg:
                print("[OCR GPU PREFLIGHT] FAIL: cudnn64_8.dll missing or incompatible", file=sys.stderr)
                print("[OCR GPU PREFLIGHT] FAIL: cudnn64_8.dll missing or incompatible")
                if not allow_cpu_fallback:
                    print("PaddleOCR GPU preflight failed because cudnn64_8.dll is missing or incompatible.", file=sys.stderr)
                    sys.exit(5)
            else:
                print(f"[OCR GPU PREFLIGHT] FAIL: {e}", file=sys.stderr)
                print(f"[OCR GPU PREFLIGHT] FAIL: {e}")
                if not allow_cpu_fallback:
                    print(f"PaddleOCR GPU preflight failed: {e}", file=sys.stderr)
                    sys.exit(5)

            # Handle CPU fallback if allowed
            if allow_cpu_fallback:
                print("Warning: OCR GPU preflight failed. Falling back to CPU.")
                if "device" in sig.parameters:
                    ocr_params["device"] = "cpu"
                else:
                    ocr_params["use_gpu"] = False
                    if "gpu_mem" in ocr_params:
                        del ocr_params["gpu_mem"]
                try:
                    reader = PaddleOCR(**ocr_params)
                    run_on_gpu = False
                except Exception as ex:
                    print(f"PaddleOCR CPU fallback initialization failed.\nReason: {ex}", file=sys.stderr)
                    sys.exit(4)

    if run_on_gpu:
        print("[OCR GPU QUALITY PREFLIGHT] running Chinese text probe...")
        try:
            quality_ok, quality_details = _run_ocr_quality_probe(reader)
            print("[OCR GPU QUALITY PREFLIGHT] " + json.dumps(quality_details, ensure_ascii=False))
        except Exception as e:
            quality_ok = False
            quality_details = {"ok": False, "error": str(e)}
            print(f"[OCR GPU QUALITY PREFLIGHT] FAIL: {e}", file=sys.stderr)

        if quality_ok:
            print("[OCR GPU QUALITY PREFLIGHT] PASS")
        elif not allow_cpu_fallback:
            print("PaddleOCR GPU quality preflight failed. GPU OCR is returning invalid Chinese text.", file=sys.stderr)
            print(json.dumps(quality_details, ensure_ascii=False), file=sys.stderr)
            sys.exit(5)
        else:
            print("Warning: OCR GPU quality preflight failed. Falling back to CPU.")
            if "device" in sig.parameters:
                ocr_params["device"] = "cpu"
            else:
                ocr_params["use_gpu"] = False
                if "gpu_mem" in ocr_params:
                    del ocr_params["gpu_mem"]
            try:
                reader = PaddleOCR(**ocr_params)
                cpu_ok, cpu_details = _run_ocr_quality_probe(reader)
                print("[OCR CPU QUALITY PREFLIGHT] " + json.dumps(cpu_details, ensure_ascii=False))
                if not cpu_ok:
                    print("PaddleOCR CPU quality preflight also failed.", file=sys.stderr)
                    sys.exit(5)
                run_on_gpu = False
                print("[OCR CPU QUALITY PREFLIGHT] PASS")
            except SystemExit:
                raise
            except Exception as ex:
                print(f"PaddleOCR CPU fallback quality preflight failed.\nReason: {ex}", file=sys.stderr)
                sys.exit(5)

    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        print(f"Error: Cannot open video file {video_path}", file=sys.stderr)
        sys.exit(2)

    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = total_frames / fps if fps > 0 else 0
    vid_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    vid_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    print(f"Video resolution: {vid_w}x{vid_h} | FPS: {fps:.2f} | Duration: {duration:.2f}s")

    if args.region:
        try:
            rx, ry, rw, rh, orig_w, orig_h = map(int, args.region.split(','))

            norm_x = rx / orig_w if orig_w > 0 else 0.0
            norm_y = ry / orig_h if orig_h > 0 else 0.0
            norm_w = rw / orig_w if orig_w > 0 else 0.0
            norm_h = rh / orig_h if orig_h > 0 else 0.0

            real_x = int(round(norm_x * vid_w))
            real_y = int(round(norm_y * vid_h))
            real_w = max(1, int(round(norm_w * vid_w)))
            real_h = max(1, int(round(norm_h * vid_h)))

            real_w = min(real_w, vid_w)
            real_h = min(real_h, vid_h)
            real_x = max(0, min(real_x, vid_w - real_w))
            real_y = max(0, min(real_y, vid_h - real_h))

            print(f"[OCR] [UI Cords] region={rx},{ry},{rw},{rh}@{orig_w}x{orig_h}")
            print(f"[OCR] [Normalized Cords] x={norm_x:.4f}, y={norm_y:.4f}, w={norm_w:.4f}, h={norm_h:.4f}")
            print(f"[OCR] [Native Frame Cords] x={real_x}, y={real_y}, w={real_w}, h={real_h} (frame {vid_w}x{vid_h})")
        except Exception as e:
            print(f"Error parsing region '{args.region}': {e}", file=sys.stderr)
            sys.exit(1)
    else:
        # Default to bottom 25% of the frame
        real_x = 0
        real_y = int(vid_h * 0.75)
        real_w = vid_w
        real_h = int(vid_h * 0.25)
        print(f"OCR Region: default bottom-center (x={real_x}, y={real_y}, w={real_w}, h={real_h})")

    sample_interval_sec = args.sample_interval
    frame_interval = max(1, int(fps * sample_interval_sec))

    print(f"Analyzing video (sampling every {sample_interval_sec}s)...")

    raw_detections = []
    prev_gray = None
    last_text = ""
    last_raw_text = ""
    last_confidence = 0.0
    last_reject_reason = None
    ocr_calls = 0
    skipped_calls = 0 # Duplicate-frame skips
    ocr_time_sum = 0.0
    frame_idx = 0
    start_time = time.time()

    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                break
        
            if frame_idx % frame_interval == 0:
                current_time_sec = frame_idx / fps
                crop = frame[real_y:real_y+real_h, real_x:real_x+real_w]
                
                if crop.size == 0:
                    frame_idx += 1
                    continue
        
                gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
                is_duplicate = False
                
                if prev_gray is not None:
                    diff = cv2.absdiff(gray, prev_gray)
                    mean_diff = diff.mean()
                    if mean_diff < 2.0:
                        is_duplicate = True
                        skipped_calls += 1
        
                if is_duplicate:
                    text = last_text
                    raw_text = last_raw_text
                    avg_confidence = last_confidence
                    reject_reason = last_reject_reason
                else:
                    ocr_start = time.time()
                    results = reader.ocr(gray, cls=True)
                    ocr_time_sum += (time.time() - ocr_start)
                    ocr_calls += 1
                    
                    text_list = []
                    scores = []
                    if results and results[0]:
                        for detected_text, score in _extract_ocr_lines(results, args.min_line_confidence):
                            text_list.append(detected_text)
                            scores.append(score)
                    raw_text = " ".join(text_list).strip()
                    raw_text = " ".join(raw_text.split())
                    avg_confidence = sum(scores) / len(scores) if scores else 0.0
                    reject_reason = _ocr_garbage_reason(raw_text, args.max_frame_text_chars)
                    text = "" if reject_reason else raw_text
                    last_text = text
                    last_raw_text = raw_text
                    last_confidence = avg_confidence
                    last_reject_reason = reject_reason
                    prev_gray = gray
        
                raw_detections.append((current_time_sec, text, raw_text, avg_confidence, reject_reason))
        
                if total_frames > 0:
                    pct = min(100, int((frame_idx / total_frames) * 100))
                    print(f"PROGRESS: {pct}%", flush=True)
        
            frame_idx += 1
    except Exception as e:
        print(f"OCR processing failed during scanning: {e}", file=sys.stderr)
        traceback.print_exc()
        sys.exit(5)

    cap.release()
    print("PROGRESS: 100%", flush=True)
    elapsed = time.time() - start_time
    
    # Total frames skipped from OCR includes both the interval skips and duplicate skips
    skipped_frames_total = (frame_idx - ocr_calls)
    avg_ocr_time = (ocr_time_sum / ocr_calls) if ocr_calls > 0 else 0.0

    print(f"Scanning completed in {elapsed:.1f}s.")
    print(f"OCR calls: {ocr_calls}")
    print(f"Skipped frames: {skipped_frames_total}")
    print(f"Duplicate-frame skips: {skipped_calls}")
    print(f"Average OCR time per call: {avg_ocr_time:.3f}s")

    accepted_frames = sum(1 for _t, txt, _raw, _conf, _reason in raw_detections if txt)
    raw_text_frames = sum(1 for _t, _txt, raw, _conf, _reason in raw_detections if raw)
    rejected_frames = sum(1 for _t, _txt, _raw, _conf, reason in raw_detections if reason)

    # Blocker 5: Save raw detections to ocr_raw.json
    raw_json_path = output_srt_path.parent / "ocr_raw.json"
    raw_payload = [
        {
            "timestamp": t,
            "text": txt,
            "raw_text": raw_txt,
            "confidence": round(float(conf), 4),
            "rejected": bool(reason),
            "reject_reason": reason,
        }
        for t, txt, raw_txt, conf, reason in raw_detections
    ]
    try:
        _atomic_write_json(raw_json_path, raw_payload)
        print(f"Saved raw detections to {raw_json_path}")
    except Exception as e:
        print(f"Warning: Failed to save ocr_raw.json: {e}", file=sys.stderr)

    # Temporal Smoothing & Merge
    srt_blocks = []
    block_id = 1
    active_text = ""
    active_start = -1.0
    active_end = -1.0
    
    gap_threshold = 0.4
    similarity_threshold = 0.7
    min_duration = 0.3

    try:
        for current_time, text, _raw_text, _confidence, _reject_reason in raw_detections:
            if len(text) < 1:
                text = ""

            if text:
                if not active_text:
                    active_text = text
                    active_start = current_time
                    active_end = current_time + sample_interval_sec
                else:
                    sim = get_similarity(active_text, text)
                    forced_split = (
                        float(args.max_block_duration) > 0
                        and (current_time + sample_interval_sec - active_start) > float(args.max_block_duration)
                    )
                    is_same_subtitle = (
                        sim >= similarity_threshold
                        or (text in active_text and len(text) >= 2)
                        or (active_text in text and len(active_text) >= 2)
                    )
                    if is_same_subtitle and not forced_split:
                        active_end = current_time + sample_interval_sec
                        if len(text) > len(active_text):
                            active_text = text
                    else:
                        dur = active_end - active_start
                        if dur >= min_duration:
                            srt_blocks.append({
                                "id": block_id,
                                "start": active_start,
                                "end": active_end,
                                "text": active_text
                            })
                            block_id += 1
                        active_text = text
                        active_start = current_time
                        active_end = current_time + sample_interval_sec
            else:
                if active_text:
                    if current_time - active_end <= gap_threshold:
                        continue
                    else:
                        dur = active_end - active_start
                        if dur >= min_duration:
                            srt_blocks.append({
                                "id": block_id,
                                "start": active_start,
                                "end": active_end,
                                "text": active_text
                            })
                            block_id += 1
                        active_text = ""
                        active_start = -1.0
                        active_end = -1.0

        if active_text:
            dur = active_end - active_start
            if dur >= min_duration:
                srt_blocks.append({
                    "id": block_id,
                    "start": active_start,
                    "end": active_end,
                    "text": active_text
                })
    except Exception as e:
        print(f"OCR processing failed during smoothing: {e}", file=sys.stderr)
        traceback.print_exc()
        sys.exit(5)
    
    # Blocker 5: Save srt_blocks to ocr_grouped_blocks.json
    grouped_json_path = output_srt_path.parent / "ocr_grouped_blocks.json"
    grouped_payload = [
        {"id": b["id"], "start": b["start"], "end": b["end"], "text": b["text"]}
        for b in srt_blocks
    ]
    try:
        _atomic_write_json(grouped_json_path, grouped_payload)
        print(f"Saved grouped blocks to {grouped_json_path}")
    except Exception as e:
        print(f"Warning: Failed to save ocr_grouped_blocks.json: {e}", file=sys.stderr)

    summary_path = output_srt_path.parent / "ocr_summary.json"
    summary_payload = {
        "status": "OK" if srt_blocks else "EMPTY",
        "device": "gpu" if run_on_gpu else "cpu",
        "duration_sec": round(float(duration), 3),
        "sample_interval_sec": sample_interval_sec,
        "ocr_calls": ocr_calls,
        "raw_detection_frames": len(raw_detections),
        "raw_text_frames": raw_text_frames,
        "accepted_frames": accepted_frames,
        "rejected_frames": rejected_frames,
        "srt_blocks": len(srt_blocks),
        "region": {
            "x": real_x,
            "y": real_y,
            "w": real_w,
            "h": real_h,
            "video_w": vid_w,
            "video_h": vid_h,
        },
    }
    try:
        _atomic_write_json(summary_path, summary_payload)
        print(f"Saved OCR summary to {summary_path}")
    except Exception as e:
        print(f"Warning: Failed to save ocr_summary.json: {e}", file=sys.stderr)

    # Output SRT file
    srt_content = []
    for block in srt_blocks:
        srt_content.append(f"{block['id']}\n")
        srt_content.append(f"{format_srt_time(block['start'])} --> {format_srt_time(block['end'])}\n")
        srt_content.append(f"{block['text']}\n\n")
    _atomic_write_text(output_srt_path, "".join(srt_content))

    if not srt_blocks:
        print(
            "[OCR EMPTY] No valid subtitle blocks were produced. "
            f"raw_text_frames={raw_text_frames}, accepted_frames={accepted_frames}, rejected_frames={rejected_frames}. "
            "Check OCR region, GPU OCR quality, and ocr_raw.json/ocr_summary.json.",
            file=sys.stderr,
        )
        sys.exit(6)

    print(f"OCR transcription completed successfully. Saved {len(srt_blocks)} blocks to {output_srt_path}")

    # Cleanup VRAM and exit
    try:
        del reader
    except NameError:
        pass
    import gc
    gc.collect()
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except Exception:
        pass
    sys.exit(0)

if __name__ == "__main__":
    main()