import pathlib
import sys
import time
from dataclasses import dataclass

import click
import numpy as np
import PIL.Image


@dataclass(slots=True)
class FileInfo:
    shape: (int, int)
    filesize: int

    @property
    def n_pixels(self):
        return self.shape[0] * self.shape[1]


def _sort_items(items, spec):
    def f(item, c):
        match c:
            case "p":
                return item[0]
            case "i":
                return item[1].n_pixels
            case "f":
                return item[1].filesize
            case _:
                raise Exception(f"unrecognized sort char: {c} ")

    if spec:
        return sorted(items, key=lambda item: [f(item, c) for c in spec])
    else:
        return items


@click.command()
@click.option("-s", "--sort")
def main(sort: str | None) -> None:
    d = dict()
    errors = dict()
    ctr = np.int64(0)

    start_time = time.perf_counter()
    for line in sys.stdin:
        path = pathlib.Path(line.strip())
        try:
            filesize = path.stat().st_size
            if path.suffix == ".webp":
                d[str(path)] = FileInfo(webp_dims(path), filesize)
            else:
                im = PIL.Image.open(path)
                # _exif = im.getexif()
                d[str(path)] = FileInfo(im.size, filesize)
            ctr += 1
        except Exception as e:
            errors[path] = e
    scan_time = time.perf_counter() - start_time

    path_width = max(map(len, d))
    for path, fi in _sort_items(d.items(), sort):
        fsize = int(fi.filesize / 1024)
        mpix = fi.n_pixels / 1_000_000
        w, h = fi.shape
        print(f"{path:{path_width}} {w:4} {h:4} {mpix:4.1f} Mpx {fsize:5d} kiB")

    print(f"{ctr} files scanned in {scan_time}s")
    print(f"{len(errors)} Files failed to open. They may not be image files")


def webp_dims(path: str) -> (int, int):
    "Extract the canvas size of a webp image"
    MAGIC = b"\x9d\x01\x2a"
    MASK = 0x3FFF
    with open(path, "rb") as fp:
        data = fp.read(0x400)  # Assumes VP8 headers is in first 1024 bytes
        if data[12:16] == b"VP8X":
            w = int.from_bytes(data[24:27], "little") + 1
            h = int.from_bytes(data[27:30], "little") + 1
        else:
            off = data.find(MAGIC)
            w = int.from_bytes(data[off + 3 : off + 5], "little") & MASK
            h = int.from_bytes(data[off + 5 : off + 7], "little") & MASK
        return w, h
