playground/scan_vids.py

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "psycopg[binary]",
# ]
# ///
import argparse
import contextlib
import hashlib
import json
import logging
import subprocess
import typing
from pathlib import Path

import psycopg
import psycopg.rows

import ffmpeg


def hash_partial(f: Path) -> str:
    sha1 = hashlib.sha1()
    chunk_size = 1024 * 1024 * 10  # 10MB chunk size

    total_read = 0
    with f.open("rb") as file:
        while chunk := file.read(chunk_size):
            total_read += len(chunk)
            sha1.update(chunk)
            break  # Only reads the first 10MB

    return f"sha1:{total_read}:{sha1.hexdigest()}"


@contextlib.contextmanager
def connect_db() -> typing.Generator[psycopg.Connection, typing.Any, typing.Any]:
    with psycopg.connect("postgres://abdus:abdus@db.abdus.dev:5444/smut?sslmode=disable") as conn:
        conn.row_factory = psycopg.rows.dict_row
        yield conn


def video_exists(con: psycopg.Connection, file_path: str) -> bool:
    stmt = con.execute(
        "UPDATE videos SET last_seen_at = NOW() WHERE file_path = %s RETURNING true AS exists",
        [file_path],
    )
    row = stmt.fetchone()
    return row["exists"]


def delete_video(con: psycopg.Connection, file_path: str) -> bool:
    con.execute("DELETE FROM videos WHERE file_path = %s", [file_path])


def save_video(
    con: psycopg.Connection,
    category: str,
    file_path: str,
    ffprobe: dict,
    hash_partial: str,
):
    sql = """INSERT INTO videos(category, file_path, ffprobe, hash_partial) VALUES (%s, %s, %s, %s) ON CONFLICT(file_path) SET last_seen_at = NOW();"""
    con.execute(
        sql,
        [
            category,
            file_path,
            json.dumps(ffprobe),
            hash_partial,
        ],
    )


def save_error(con: psycopg.Connection, file_path: str, stderr: str):
    sql = """INSERT INTO scan_state(file_path, stderr) VALUES (%s, %s) ON conflict do nothing"""
    con.execute(
        sql,
        [
            file_path,
            stderr,
        ],
    )


def scan_videos(con: psycopg.Connection, video_paths: list[Path]) -> None:
    for i, f in enumerate(video_paths):
        if "/_picks/" in str(f):
            continue
        progress = f"[{i + 1}/{len(video_paths)}]"
        if f.is_dir():
            continue

        file_path = str(f.absolute())

        if video_exists(con, file_path=file_path):
            logging.debug(f"{progress} video already exists {f=}")
            continue

        logging.debug(f"probing {f=}")
        try:
            probed = ffmpeg.ffprobe(f)
        except subprocess.CalledProcessError as e:
            save_error(con, file_path=file_path, stderr=e.stderr)
            continue

        file_hash = hash_partial(f)

        save_video(
            category=category_for_file(f),
            con=con,
            file_path=file_path,
            ffprobe=probed.__dict__,
            hash_partial=file_hash,
        )
        logging.info(f"{progress} saved {f=}")
        con.commit()


def category_for_file(file_path: Path) -> str:
    if str(file_path).startswith("/Volumes/"):
        return f"downloaded.{file_path.parts[1].lower()}"
    if str(file_path).startswith("/mnt/box"):
        return "box"
    raise ValueError(f"unknown category for {file_path}")


def find_copies(con: psycopg.Connection, video_path: Path) -> list[str]:
    probed = ffmpeg.ffprobe(video_path)

    sql = """
SELECT
    file_path
FROM videos
WHERE
    file_path != %(file_path)s
    AND ( 
        hash_partial = %(hash_partial)s
        OR (abs(duration_sec - %(duration_sec)s) < 0.1 AND abs(size_bytes - %(size_bytes)s) < 3000000)
    );
   """

    stmt: psycopg.Cursor = con.execute(
        sql,
        {
            "file_path": str(video_path.absolute()),
            "hash_partial": hash_partial(video_path),
            "duration_sec": probed.duration_sec,
            "size_bytes": probed.size_bytes,
        },
    )
    rows = stmt.fetchall()

    return [it["file_path"] for it in rows]


def parse_args():
    arger = argparse.ArgumentParser()
    arger.add_argument("video_paths", nargs="+", type=Path, help="Path to the videos")
    arger.add_argument("--debug", action="store_true", help="Path to the videos")

    return arger.parse_args()


def main():
    logging.basicConfig(format=logging.BASIC_FORMAT, level=logging.INFO)
    args = parse_args()
    if args.debug:
        logging.getLogger().setLevel(logging.DEBUG)

    video_paths: list[Path] = args.video_paths
    if not video_paths:
        logging.info("no files to scan")
        return

    with connect_db() as con:
        scan_videos(con=con, video_paths=video_paths)


if __name__ == "__main__":
    main()