main.py

import platform
import shutil
import re
from datetime import datetime, timedelta
import cv2
import sys
import numpy as np
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path
from collections import deque
import statistics
import pytesseract
# Set Tesseract path based on OS
if sys.platform.startswith("win"):
    # Default Windows path
    tesseract_default = r"C:\Program Files\Tesseract-OCR\tesseract.exe"
    if Path(tesseract_default).exists():
        pytesseract.pytesseract.tesseract_cmd = tesseract_default
    else:
        raise FileNotFoundError(
            f"Tesseract executable not found at {tesseract_default}. "
            "Please install Tesseract OCR or adjust the path."
        )
else:
    # Unix / Linux / macOS default
    tesseract_default = "/usr/bin/tesseract"
    if Path(tesseract_default).exists():
        pytesseract.pytesseract.tesseract_cmd = tesseract_default
    else:
        raise FileNotFoundError(
            f"Tesseract executable not found at {tesseract_default}. "
            "Please install Tesseract OCR or adjust the path."
        )

import typer

from process_frames import process


# ============================================================
#  SYSTEM CHECKS
# ============================================================

def check_dependency_exists(name: str) -> bool:
    """Check if a binary exists in PATH."""
    return shutil.which(name) is not None


def check_system_dependencies():
    print("\n[ Checking system dependencies ]")

    # --- Check FFmpeg ---
    if not check_dependency_exists("ffmpeg"):
        print("⚠ WARNING: FFmpeg is not installed or not found in PATH.")
        print("   → Install from: https://ffmpeg.org/download.html")
    else:
        print("✔ FFmpeg found")

    # --- Check Tesseract ---
    if pytesseract.pytesseract.tesseract_cmd == '':
        print("⚠ WARNING: Tesseract OCR is not installed or not found in PATH.")
        print("   → Install from: https://tesseract-ocr.github.io/tessdoc/")
    else:
        print("✔ Tesseract found")

    print("")


# ============================================================
#  CONFIGURE TESSERACT BASED ON OS
# ============================================================

def configure_tesseract():
    system = platform.system().lower()

    if "windows" in system:
        default_path = r"C:\Program Files\Tesseract-OCR\tesseract.exe"
        if Path(default_path).exists():
            pytesseract.pytesseract.tesseract_cmd = default_path
        else:
            print("⚠ Warning: Default Windows Tesseract path not found.")
            print("   Using system PATH fallback.")
    else:
        # Linux / Mac usually use PATH tesseract
        pytesseract.pytesseract.tesseract_cmd = '/usr/bin/tesseract'


# ============================================================
#  OCR CLEANING & REPAIR
# ============================================================

OCR_REPLACEMENTS = {
    '—': '-', '_': '-', '.': ':', ',': ':',
    'O': '0', 'o': '0', 'I': '1', 'l': '1', 'L': '1',
    'S': '5', 's': '5', 'B': '8', 'Z': '2', 'z': '2',
    'g': '9', 'q': '9', '|': '1', '¦': '1', '¡': '1',
    'f': '1', '/': ':', '\\': ':', '‘': "'", '’': "'",
}
_translate_table = str.maketrans(OCR_REPLACEMENTS)


def fuzzy_clean(s: str) -> str:
    if not s:
        return ""
    s = s.strip().replace('\r', ' ').replace('\n', ' ')
    s = s.translate(_translate_table)
    s = re.sub(r'[^0-9A-Za-z:\-_/T\s]', '', s)
    s = re.sub(r'\s+', ' ', s).strip()
    return s


TIMESTAMP_PATTERNS = [
    r'(?P<year>20\d{2})[-_/]?(?P<mon>\d{1,2})[-_/]?(?P<day>\d{1,2})[ T_,-]*'
    r'(?P<hour>\d{1,2})[:_,-]?(?P<min>\d{1,2})[:_,-]?(?P<sec>\d{1,2})'
    r'(?:[:._,-]?(?P<sub>\d{1,6}))?(?:[_\s-]*(?P<frame>\d{1,4}))?',

    r'(?P<digits>20\d{10,}\d*)'
]


def parse_digits_fallback(digits: str):
    if len(digits) < 12:
        return None
    if len(digits) >= 14:
        y = digits[0:4]
        mo = digits[4:6]
        d = digits[6:8]
        H = digits[8:10]
        M = digits[10:12]
        S = digits[12:14]
        frame = digits[14:] or None
    else:
        y = digits[0:4]
        mo = digits[4:6] if len(digits) >= 6 else '01'
        d = digits[6:8] if len(digits) >= 8 else '01'
        H = digits[8:10] if len(digits) >= 10 else '00'
        M = digits[10:12] if len(digits) >= 12 else '00'
        S = '00'
        frame = digits[12:] or None

    try:
        return {
            'year': int(y), 'mon': int(mo), 'day': int(d),
            'hour': int(H), 'min': int(M), 'sec': int(S),
            'frame': frame
        }
    except:
        return None


def repair_ocr_timestamp(text: str, video_start_year=None):
    if not text:
        return None, None

    text = fuzzy_clean(text)

    for patt in TIMESTAMP_PATTERNS:
        m = re.search(patt, text)
        if not m:
            continue
        gd = m.groupdict()

        if 'digits' in gd and gd['digits']:
            parsed = parse_digits_fallback(gd['digits'])
            if not parsed:
                continue
        else:
            try:
                parsed = {
                    'year': int(gd['year']),
                    'mon': int(gd['mon'] or 1),
                    'day': int(gd['day'] or 1),
                    'hour': int(gd['hour'] or 0),
                    'min': int(gd['min'] or 0),
                    'sec': int(gd['sec'] or 0),
                    'frame': gd.get('frame')
                }
            except:
                continue

        year = parsed['year']
        if not (2000 <= year <= 2100) and video_start_year:
            year = video_start_year

        try:
            dt = datetime(
                year,
                max(1, min(12, parsed['mon'])),
                max(1, min(31, parsed['day'])),
                max(0, min(23, parsed['hour'])),
                max(0, min(59, parsed['min'])),
                max(0, min(59, parsed['sec'])),
            )
            return dt, parsed.get('frame')
        except:
            continue

    # fallback
    digits = re.sub(r'\D', '', text)
    fallback = parse_digits_fallback(digits) if digits else None
    if fallback:
        year = fallback['year']
        if not (2000 <= year <= 2100) and video_start_year:
            year = video_start_year
        try:
            dt = datetime(
                year,
                max(1, min(12, fallback['mon'])),

                max(1, min(28, fallback['day'])),
                max(0, min(23, fallback['hour'])),
                max(0, min(59, fallback['min'])),
                max(0, min(59, fallback['sec'])),
            )
            return dt, fallback.get('frame')
        except:
            pass

    return None, None


# ============================================================
#  TIMESTAMP SMOOTHER
# ============================================================

class TimestampSmoother:
    def __init__(self, window=120):
        self.last_ts = None
        self.deltas = deque(maxlen=window)

    def accept(self, candidate: datetime) -> bool:
        if candidate is None:
            return False
        if self.last_ts is None:
            self.last_ts = candidate
            return True
        if candidate < self.last_ts:
            return False

        delta = (candidate - self.last_ts).total_seconds()

        if delta == 0:
            self.deltas.append(0.0)
            self.last_ts = candidate
            return True

        positive = [d for d in self.deltas if d > 0]
        median_delta = (
            statistics.median(positive) if len(positive) >= 3 else
            statistics.mean(positive) if positive else None
        )

        if median_delta is None:
            if delta > 10:
                return False
            self.deltas.append(delta)
            self.last_ts = candidate
            return True

        if delta <= max(median_delta * 3, 2.0):
            self.deltas.append(delta)
            self.last_ts = candidate
            return True

        return False

    def force_accept(self, candidate: datetime):
        if candidate is None:
            return
        if self.last_ts is None:
            self.last_ts = candidate
            return
        delta = (candidate - self.last_ts).total_seconds()
        if delta < 0:
            return
        self.deltas.append(max(0.0, delta))
        self.last_ts = candidate

    def expected_next(self) -> datetime:
        if self.last_ts is None:
            return None
        positive = [d for d in self.deltas if d > 0]
        if not positive:
            return self.last_ts
        median_delta = statistics.median(positive)
        if median_delta <= 0:
            return self.last_ts
        return self.last_ts + timedelta(seconds=median_delta)


# ============================================================
#  IMAGE EXTRACTION
# ============================================================

def remove_noise_gaussian(image):
    return cv2.GaussianBlur(image, (5, 5), 0)


def extract_green_text(image):
    b, g, r = cv2.split(image)
    mask = (g > 150) & (r < 80) & (b < 80)
    result = np.zeros_like(g, dtype=np.uint8)
    result[mask] = 255
    kernel = np.ones((3, 3), np.uint8)
    return cv2.erode(cv2.dilate(result, kernel, iterations=1), kernel, iterations=1)


def crop_to_green_text(image):
    mask = extract_green_text(image)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if contours:
        xs = [cv2.boundingRect(c)[0] for c in contours]
        ys = [cv2.boundingRect(c)[1] for c in contours]
        ws = [cv2.boundingRect(c)[0] + cv2.boundingRect(c)[2] for c in contours]
        hs = [cv2.boundingRect(c)[1] + cv2.boundingRect(c)[3] for c in contours]

        x, y = max(0, min(xs) - 6), max(0, min(ys) - 6)
        w = min(image.shape[1] - x, max(ws) - min(xs) + 12)
        h = min(image.shape[0] - y, max(hs) - min(ys) + 12)
        return mask[y:y + h, x:x + w]
    return None


# ============================================================
#  FILENAME TIMESTAMP PARSE
# ============================================================

VIDEO_TS_REGEX = re.compile(r'(20\d{6}T\d{6})')


def parse_video_filename_start_ts(name):
    m = VIDEO_TS_REGEX.search(name)
    if not m:
        return None
    try:
        return datetime.strptime(m.group(1), "%Y%m%dT%H%M%S")
    except:
        return None


# ============================================================
# PROCESS SINGLE VIDEO
# ============================================================

CONSENSUS_REQUIRED = 3


def process_single_video(video_file: Path):
    print(f"▶ Starting video: {video_file}")
    cap = cv2.VideoCapture(str(video_file))
    if not cap.isOpened():
        print(f"Cannot open: {video_file}")
        return

    smoother = TimestampSmoother()
    frame_counter_global = {}
    video_start_ts = parse_video_filename_start_ts(video_file.name)

    frame_count = 0
    consensus_last = None
    consensus_count = 0

    frames_dir = video_file.parent / "frames"
    frames_dir.mkdir(parents=True, exist_ok=True)

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        frame_count += 1

        if frame_count % 2 == 0:
            print(f"   progress {frame_count}/{total_frames} frames")

        cropped = crop_to_green_text(frame)
        if cropped is None:
            continue

        if cropped.shape[1] >= 470:
            cropped = cropped[:, 60:]
        if cropped.shape[0] >= 50:
            cropped = cropped[50:, :]

        if cropped is None or cropped.size == 0:
            continue
        if cropped.shape[0] < 5 or cropped.shape[1] < 5:
            continue

        gray = cv2.bitwise_not(cropped)

        cropped = cv2.adaptiveThreshold(
            gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv2.THRESH_BINARY, 11, 2
        )
        cropped = remove_noise_gaussian(cropped)

        raw_text = pytesseract.image_to_string(cropped, config='--psm 7')
        dt_candidate, frame_sub = repair_ocr_timestamp(
            raw_text,
            video_start_year=video_start_ts.year if video_start_ts else None
        )

        accepted = dt_candidate is not None and smoother.accept(dt_candidate)

        if not accepted and dt_candidate is not None:
            if consensus_last == dt_candidate:
                consensus_count += 1
            else:
                consensus_last = dt_candidate
                consensus_count = 1

            if consensus_count >= CONSENSUS_REQUIRED:
                smoother.force_accept(dt_candidate)
                accepted = True
                consensus_last = None
                consensus_count = 0
        else:
            consensus_last = None
            consensus_count = 0

        if not accepted:
            est = smoother.expected_next()
            dt_to_use = est or datetime.utcnow()
            debug = "(ESTIMATED)"
        else:
            dt_to_use = dt_candidate
            debug = "(OCR)"

        ts_base = dt_to_use.strftime("%Y_%m_%dT%H_%M_%S")
        if frame_sub:
            ts_base += f"_{str(frame_sub).zfill(3)}"

        count = frame_counter_global.get(ts_base, 0) + 1
        frame_counter_global[ts_base] = count

        filename = f"{ts_base}_{count:03}.jpg"
        filepath = frames_dir / filename

        cv2.imwrite(str(filepath), frame, [int(cv2.IMWRITE_JPEG_QUALITY), 70])

        print(f"{frame_count:6d} {debug} -> {dt_to_use}  file: {filename}  raw: {raw_text.strip()}")

    cap.release()
    print(f"✔ Finished video: {video_file.name} ({frame_count} frames)\n")


# ============================================================
# MAIN PROCESSING FUNCTION (TYPER ENTRYPOINT)
# ============================================================

app = typer.Typer(add_completion=False)


default_dirs = [
    Path("C:/sync/24/2025Sep12/videos"),
    Path("C:/sync/25/2025Sep12/videos"),
    Path("C:/sync/26/2025Sep12/videos"),
    Path("C:/sync/27/2025Sep12/videos"),
]


@app.command()
def run(
    video_dirs: list[Path] = typer.Option(
        default_dirs,
        help="List of directories containing .mp4 videos. Defaults to the 4 sync paths."
    ),
    max_workers: int = typer.Option(
        12,
        help="Number of parallel processes (default 12)."
    ),
):
    """
    Process videos from directories, extract OCR timestamps, and save frames.

    If no directories are supplied, defaults to:
    default_dirs
    """
    print("\n=== VIDEO OCR FRAME EXTRACTION ===\n")

    # --- Check environment ---
    configure_tesseract()
    check_system_dependencies()

    # collect ALL mp4 files
    all_videos = []
    for vd in video_dirs:
        vids = list(vd.glob("*.mp4"))
        all_videos.extend(vids)
        print(f"Found {len(vids)} videos in {vd}")

    if not all_videos:
        print("No videos found. Exiting.")
        raise typer.Exit(1)

    print(f"\n✔ Total videos: {len(all_videos)}\n")

    # Run each video in parallel
    with ProcessPoolExecutor(max_workers=max_workers) as ex:
        for i, _ in enumerate(ex.map(process_single_video, all_videos), start=1):
            print(f"[Progress] Processed {i}/{len(all_videos)} videos\n")

    print("\n=== Processing OCR Post-Step ===")
    process()
    print("=== DONE ===\n")


# ============================================================
# ENTRY POINT
# ============================================================

if __name__ == "__main__":
    app()