#!/usr/bin/env python3
"""Correct blob sizes in a blossom-server sqlite.db from the files on disk.

blossom-server 4.4.1 sometimes recorded a wrong size at upload (0, or a
partial count — fixed upstream in 4.6.0). 4.x ignored the stored size when
serving, but 6.x sends it as Content-Length, so affected blobs come back
empty or truncated after the 4.x -> 6.x migration. Only rows whose file
hashes to its sha256 are touched, so a corrupt file is reported, not masked.

Run with blossom stopped, against the volume's data dir:
    python3 scripts/fix-blossom-sizes.py <blossom-data-dir> [--apply]
Without --apply it only reports.
"""
import hashlib
import os
import sqlite3
import sys


def sha256_file(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def main() -> int:
    args = [a for a in sys.argv[1:] if a != "--apply"]
    apply = "--apply" in sys.argv
    if len(args) != 1:
        print(__doc__)
        return 2
    data_dir = args[0]
    blobs_dir = os.path.join(data_dir, "blobs")
    files = {name[:64]: name for name in os.listdir(blobs_dir)}

    db = sqlite3.connect(os.path.join(data_dir, "sqlite.db"))
    fixed = problems = 0
    for sha256, size in db.execute("SELECT sha256, size FROM blobs").fetchall():
        name = files.get(sha256)
        if name is None:
            print(f"MISSING  {sha256}: no file on disk")
            problems += 1
            continue
        path = os.path.join(blobs_dir, name)
        real = os.path.getsize(path)
        if real == size:
            continue
        if sha256_file(path) != sha256:
            print(f"CORRUPT  {name}: file does not hash to its name; size left at {size}")
            problems += 1
            continue
        print(f"{'FIXED' if apply else 'WOULD FIX'}  {name}: {size} -> {real}")
        if apply:
            db.execute("UPDATE blobs SET size = ? WHERE sha256 = ?", (real, sha256))
        fixed += 1
    db.commit()
    db.close()
    print(f"{fixed} size(s) {'corrected' if apply else 'to correct'}, {problems} problem(s)")
    return 1 if problems else 0


if __name__ == "__main__":
    sys.exit(main())