#!/usr/bin/env python3
"""Inventory a CDN tree on STARGATE.

The output is TSV on stdout:

path, size, sha256, mime, width, height

It uses only the Python standard library so it can run on STARGATE without
installing ImageMagick or other packages.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import mimetypes
import os
import struct
import sys
from pathlib import Path


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as fh:
        for chunk in iter(lambda: fh.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def png_size(path: Path) -> tuple[str, str]:
    with path.open("rb") as fh:
        data = fh.read(24)
    if data[:8] != b"\x89PNG\r\n\x1a\n" or data[12:16] != b"IHDR":
        return "", ""
    width, height = struct.unpack(">II", data[16:24])
    return str(width), str(height)


def gif_size(path: Path) -> tuple[str, str]:
    with path.open("rb") as fh:
        data = fh.read(10)
    if data[:6] not in {b"GIF87a", b"GIF89a"}:
        return "", ""
    width, height = struct.unpack("<HH", data[6:10])
    return str(width), str(height)


def jpeg_size(path: Path) -> tuple[str, str]:
    with path.open("rb") as fh:
        if fh.read(2) != b"\xff\xd8":
            return "", ""
        while True:
            marker_start = fh.read(1)
            if not marker_start:
                return "", ""
            if marker_start != b"\xff":
                continue
            marker = fh.read(1)
            while marker == b"\xff":
                marker = fh.read(1)
            if marker in {b"\xd8", b"\xd9"}:
                continue
            size_bytes = fh.read(2)
            if len(size_bytes) != 2:
                return "", ""
            segment_len = struct.unpack(">H", size_bytes)[0]
            if marker in {
                b"\xc0",
                b"\xc1",
                b"\xc2",
                b"\xc3",
                b"\xc5",
                b"\xc6",
                b"\xc7",
                b"\xc9",
                b"\xca",
                b"\xcb",
                b"\xcd",
                b"\xce",
                b"\xcf",
            }:
                data = fh.read(5)
                if len(data) != 5:
                    return "", ""
                height, width = struct.unpack(">HH", data[1:5])
                return str(width), str(height)
            fh.seek(segment_len - 2, os.SEEK_CUR)


def webp_size(path: Path) -> tuple[str, str]:
    with path.open("rb") as fh:
        data = fh.read(30)
    if data[:4] != b"RIFF" or data[8:12] != b"WEBP":
        return "", ""
    chunk = data[12:16]
    if chunk == b"VP8 " and len(data) >= 30:
        width, height = struct.unpack("<HH", data[26:30])
        return str(width & 0x3FFF), str(height & 0x3FFF)
    if chunk == b"VP8L" and len(data) >= 25:
        b0, b1, b2, b3 = data[21:25]
        width = 1 + (((b1 & 0x3F) << 8) | b0)
        height = 1 + (((b3 & 0x0F) << 10) | (b2 << 2) | ((b1 & 0xC0) >> 6))
        return str(width), str(height)
    if chunk == b"VP8X" and len(data) >= 30:
        width = 1 + int.from_bytes(data[24:27], "little")
        height = 1 + int.from_bytes(data[27:30], "little")
        return str(width), str(height)
    return "", ""


def image_size(path: Path) -> tuple[str, str]:
    suffix = path.suffix.lower()
    try:
        if suffix == ".png":
            return png_size(path)
        if suffix == ".gif":
            return gif_size(path)
        if suffix in {".jpg", ".jpeg"}:
            return jpeg_size(path)
        if suffix == ".webp":
            return webp_size(path)
    except OSError:
        return "", ""
    return "", ""


def mime_for(path: Path) -> str:
    guessed, _ = mimetypes.guess_type(path.name)
    if guessed:
        return guessed
    return ""


def iter_files(root: Path):
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames.sort()
        filenames.sort()
        for filename in filenames:
            path = Path(dirpath) / filename
            if path.is_file():
                yield path


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("root", type=Path)
    args = parser.parse_args()

    root = args.root.resolve()
    writer = csv.DictWriter(
        sys.stdout,
        fieldnames=["path", "size", "sha256", "mime", "width", "height"],
        delimiter="\t",
        lineterminator="\n",
    )
    writer.writeheader()
    for path in iter_files(root):
        rel = path.relative_to(root).as_posix()
        stat = path.stat()
        width, height = image_size(path)
        writer.writerow(
            {
                "path": rel,
                "size": stat.st_size,
                "sha256": sha256_file(path),
                "mime": mime_for(path),
                "width": width,
                "height": height,
            }
        )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())