#!/usr/bin/env python3

# FIXME use parallel hashing for many small files

import os
import sys
import sqlite3
import hashlib
import stat
import time
from concurrent.futures import ProcessPoolExecutor

sys.path.insert(0, "/nix/store/g7dr3zh6j9nmhncxw5s2ac6cg2bf6avd-python3.13-psutil-7.2.2/lib/python3.13/site-packages")
import psutil

# max_workers = os.cpu_count() # CPU threads
max_workers = psutil.cpu_count(logical=False) # CPU cores
print(f"max_workers: {max_workers}")

queue size = 1000

# no, these are not faster. the bottleneck is the HDD read speed
r'''
sys.path.insert(0, "/nix/store/ywlwgz56n9frp3xyn4vr2ifv9p93gv5c-python3.13-blake3-1.0.8/lib/python3.13/site-packages")
from blake3 import blake3

sys.path.insert(0, "/nix/store/9sn8prwd9jwf4ah368pyls7dqfp8pkhx-python3.13-xxhash-3.6.0/lib/python3.13/site-packages")
import xxhash
'''


# config

# done 340000 inodes in 150.64844965934753 seconds
# store_hashes = False

# this is really really slow...
store_hashes = True

print_progress_every_n_inodes = 10000
# print_progress_every_n_inodes = 100
# print_progress_every_n_inodes = 10

commit_db_every_n_inodes = 10000

if store_hashes:
    print_progress_every_n_inodes = 100
    commit_db_every_n_inodes = 100


# global state

done_inodes = 0
last_done_inodes = 0
start_time = time.time()
last_progress_time = start_time


SCHEMA = """
PRAGMA foreign_keys=OFF;

CREATE TABLE IF NOT EXISTS inode (
    device     INTEGER NOT NULL,
    inode      INTEGER NOT NULL,
    type       INTEGER NOT NULL,   -- 0=file,1=dir,2=other
    mode       INTEGER NOT NULL,
    size       INTEGER,
    mtime_ns   INTEGER,
    md5 BLOB,
    sha1 BLOB,
    sha256 BLOB,
    scan_id    INTEGER,
    PRIMARY KEY(device, inode)
);

CREATE TABLE IF NOT EXISTS dirent (
    parent_device INTEGER NOT NULL,
    parent_inode  INTEGER NOT NULL,
    name          TEXT NOT NULL,
    child_device  INTEGER NOT NULL,
    child_inode   INTEGER NOT NULL,
    scan_id       INTEGER,
    PRIMARY KEY(parent_device, parent_inode, name)
);

CREATE INDEX IF NOT EXISTS idx_dir_lookup
ON dirent(parent_device, parent_inode, name);

CREATE INDEX IF NOT EXISTS idx_inode_lookup
ON inode(device, inode);
"""

# 1 MiB: good for large sequential files
CHUNK = 1024 * 1024

# 2 MiB
CHUNK = 2 * 1024 * 1024

# 4 MiB: usually no real gain, sometimes worse cache locality
# CHUNK = 4 * 1024 * 1024


def compute_hashes(path):
    if not store_hashes:
        return None, None, None # test: dont hash
    md5 = hashlib.md5()
    sha1 = hashlib.sha1()
    sha256 = hashlib.sha256()
    # blake3_hasher = blake3()
    # xxh3_128 = xxhash.xxh3_128()

    with open(path, "rb") as f:
        while True:
            b = f.read(CHUNK)
            if not b:
                break
            md5.update(b)
            sha1.update(b)
            sha256.update(b)
            # blake3_hasher.update(b)
            # xxh3_128.update(b)

    return md5.digest(), sha1.digest(), sha256.digest()
    # return blake3_hasher.digest(), xxh3_128.digest(), None


def is_utf8(name: str) -> bool:
    try:
        name.encode("utf-8")
        return True
    except UnicodeEncodeError:
        return False


def scan_dir(conn, root, scan_id):
    cur = conn.cursor()

    conn.execute("UPDATE inode SET scan_id=0")
    conn.execute("UPDATE dirent SET scan_id=0")

    def maybe_commit():
        global done_inodes
        if done_inodes % commit_db_every_n_inodes == 0:
            conn.commit()
            conn.execute("BEGIN")
            # print(f"done {done_inodes} inodes", flush=True)

    def upsert_inode(dev, ino, ftype, mode, size, mtime, md5, sha1, sha256):
        global done_inodes
        global last_progress_time
        global last_done_inodes
        done_inodes += 1
        # maybe_commit()
        if done_inodes % commit_db_every_n_inodes == 0:
            conn.commit()
            conn.execute("BEGIN")
        if done_inodes % print_progress_every_n_inodes == 0:
            t2 = time.time()
            # dt = t2 - last_progress_time
            dt = t2 - start_time
            print(f"done {done_inodes} inodes in {dt} seconds = {(last_done_inodes-done_inodes)/(last_progress_time-t2)} inodes/second")
            last_progress_time = t2
            last_done_inodes = done_inodes
        cur.execute("""
        INSERT INTO inode(device, inode, type, mode, size, mtime_ns, md5, sha1, sha256, scan_id)
        VALUES(?,?,?,?,?,?,?,?,?,?)
        ON CONFLICT(device, inode) DO UPDATE SET
            type=excluded.type,
            mode=excluded.mode,
            size=excluded.size,
            mtime_ns=excluded.mtime_ns,
            md5=excluded.md5,
            sha1=excluded.sha1,
            sha256=excluded.sha256,
            scan_id=excluded.scan_id
        """, (dev, ino, ftype, mode, size, mtime, md5, sha1, sha256, scan_id))

    def upsert_dirent(pdev, pino, name, cdev, cino):
        cur.execute("""
        INSERT INTO dirent(parent_device, parent_inode, name, child_device, child_inode, scan_id)
        VALUES(?,?,?,?,?,?)
        ON CONFLICT(parent_device, parent_inode, name) DO UPDATE SET
            child_device=excluded.child_device,
            child_inode=excluded.child_inode,
            scan_id=excluded.scan_id
        """, (pdev, pino, name, cdev, cino, scan_id))

    def process_file(path, st):
        dev, ino = st.st_dev, st.st_ino

        row = cur.execute(
            "SELECT size, mtime_ns, md5, sha1, sha256 FROM inode WHERE device=? AND inode=?",
            (dev, ino),
        ).fetchone()

        md5 = sha1 = sha256 = None

        if row and row[0] == st.st_size and row[1] == st.st_mtime_ns:
            md5, sha1, sha256 = row[2], row[3], row[4]
        else:
            if stat.S_ISREG(st.st_mode):
                md5, sha1, sha256 = compute_hashes(path)

        upsert_inode(
            dev, ino,
            0 if stat.S_ISREG(st.st_mode) else 2,
            st.st_mode,
            st.st_size,
            st.st_mtime_ns,
            md5, sha1, sha256
        )

    def walk(path):
        try:
            with os.scandir(path) as it:
                entries = list(it)
        except PermissionError:
            return

        entries.sort(key=lambda e: e.name)

        try:
            pst = os.lstat(path)
            pdev, pino = pst.st_dev, pst.st_ino
        except Exception:
            return

        for e in entries:
            if not is_utf8(e.name):
                print(f"Non-UTF8 filename skipped: {e.path}", file=sys.stderr)
                continue

            try:
                if e.is_symlink():
                    st = os.lstat(e.path)
                else:
                    st = e.stat(follow_symlinks=False)
            except Exception:
                continue

            dev, ino = st.st_dev, st.st_ino

            upsert_dirent(pdev, pino, e.name, dev, ino)

            if e.is_dir(follow_symlinks=False):
                upsert_inode(dev, ino, 1, st.st_mode, None, st.st_mtime_ns, None, None, None)
                walk(e.path)
            else:
                process_file(e.path, st)

    walk(root)

    conn.execute("DELETE FROM inode WHERE scan_id != ?", (scan_id,))
    conn.execute("DELETE FROM dirent WHERE scan_id != ?", (scan_id,))

    conn.commit()


def main():
    if len(sys.argv) != 3:
        print("usage: fsindex.py ROOT DB")
        sys.exit(1)

    root = os.path.abspath(sys.argv[1])
    db = sys.argv[2]

    conn = sqlite3.connect(db)
    conn.execute("PRAGMA journal_mode=WAL;")
    conn.executescript(SCHEMA)
    conn.execute("BEGIN")

    scan_id = 1
    scan_dir(conn, root, scan_id)

    conn.commit()
    conn.close()

    # print final progress
    t2 = time.time()
    # dt = t2 - last_progress_time
    dt = t2 - start_time
    print(f"done {done_inodes} inodes in {dt} seconds = {(last_done_inodes-done_inodes)/(last_progress_time-t2)} inodes/second")


if __name__ == "__main__":
    main()