#!/usr/bin/env -S uv run --script
# /// script
# dependencies = ["psycopg[binary]"]
# ///
import dataclasses
import datetime
import hashlib
from pathlib import Path
import subprocess
import typing
import psycopg
import contextlib


@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 get_video_duration(video_path: Path) -> datetime.timedelta:
    # fmt: off
    args = [
        'ffprobe',
        '-v', 'error',
        '-show_entries', 'format=duration',
        '-of', 'default=noprint_wrappers=1:nokey=1',
        video_path,
    ]
    # fmt: on
    p = subprocess.run(args, stdout=subprocess.PIPE, check=True)
    return datetime.timedelta(seconds=round(float(p.stdout), 1))


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()}"


@dataclasses.dataclass
class SavedVideo:
    id: int
    remote_path: Path
    local_path: Path

    def __hash__(self):
        return hash(self.remote_path)


# smut=# \d videos;
#                                                                  Table "public.videos"
#          Column         |            Type             | Collation | Nullable |                                 Default
# ------------------------+-----------------------------+-----------+----------+--------------------------------------------------------------------------
#  id                     | integer                     |           | not null | nextval('videos_id_seq'::regclass)
#  category               | ltree                       |           | not null |
#  file_path              | text                        |           | not null |
#  created_at             | timestamp without time zone |           |          | now()
#  hash_partial           | text                        |           | not null |
#  ffprobe                | jsonb                       |           | not null |
#  size_bytes             | bigint                      |           | not null | generated always as ((ffprobe ->> 'size_bytes'::text)::bigint) stored
#  duration_sec           | numeric                     |           | not null | generated always as ((ffprobe ->> 'duration_sec'::text)::numeric) stored
#  suggested_filename     | text                        |           |          |
#  marked_for_deletion_at | timestamp with time zone    |           |          |
#  duration_human         | text                        |           |          | generated always as (ffprobe ->> 'duration_human'::text) stored
#  last_seen_at           | timestamp with time zone    |           |          |
# Indexes:
#     "videos_pkey" PRIMARY KEY, btree (id)
#     "videos_file_path_idx" gist (file_path gist_trgm_ops)
#     "videos_file_path_uniq" UNIQUE CONSTRAINT, btree (file_path)
def find_videos_by_hash(conn: psycopg.Connection, video_paths: list[Path]) -> list[SavedVideo]:
    file_to_hash = {p: hash_partial(p) for p in video_paths}
    if not file_to_hash:
        return []

    placeholders = ",".join(["%s"] * len(file_to_hash))
    sql = f"""
    SELECT id, file_path AS remote_path, file_path AS local_path
    FROM videos
    WHERE hash_partial IN ({placeholders})
    """
    with conn.cursor() as cur:
        rows = cur.execute(sql, list(file_to_hash.values())).fetchall()
        return [SavedVideo(**row) for row in rows]


def find_videos_by_duration(conn: psycopg.Connection, video_paths: list[Path]) -> list[SavedVideo]:
    if not video_paths:
        return []
    file_to_duration = {p: get_video_duration(p).total_seconds() for p in video_paths}
    file_to_size = {p: p.stat().st_size for p in video_paths}

    sql = f"""
    SELECT id, file_path AS remote_path, file_path AS local_path
    FROM videos
    WHERE abs(duration_sec - %s) < 0.1 AND abs(size_bytes - %s) < 1048576
    """
    with conn.cursor() as cur:
        out = []
        for file_path in video_paths:
            size = file_to_size[file_path]
            duration = file_to_duration[file_path]
            for row in cur.execute(sql, (duration, size)):
                out.append(SavedVideo(**row))
        return out


def delete_videos(con: psycopg.Connection, file_paths: list[str]) -> bool:
    with con.cursor() as cur:
        placeholders = ",".join(["%s"] * len(file_paths))
        sql = f"DELETE FROM videos WHERE file_path IN ({placeholders})"
        cur.execute(sql, file_paths)


def main():
    videos = list(Path(r"/Users/abdus/Downloads/temp").glob("*.mp4"))
    with connect_db() as conn:
        found_videos = set(find_videos_by_hash(conn, videos))
        found_videos.update(find_videos_by_duration(conn, videos))
        found_videos = {video for video in found_videos if "/mnt/box/files/_raw/prt" in video.remote_path}
        for video in found_videos:
            print(f"{video.remote_path}")
        if not found_videos:
            print("No videos found for deletion.")
            return
        input("Press Enter to continue...")

        print(f"Deleting {len(found_videos)} videos from database...")
        delete_videos(conn, [video.remote_path for video in found_videos])


if __name__ == "__main__":
    main()