"""Rewrite JPEG 2000 figures as WebP.

Twenty-one stem images extracted from the source PDFs are `.jpx`. No browser
but Safari draws JPEG 2000 — Chrome dropped it in 2015 — so those figures were
blank for almost everybody. Serving already converts them on the way out and
keeps the result, so nothing is broken by leaving them as they are; this
removes the conversion step entirely and makes what is stored the same thing
that is served.

Written as a migration of the bytes rather than a lazy rewrite because there
are three places a path is recorded — `questions.image_path`,
`questions.explanation_image_path` and `media_assets.path` — and a conversion
that updated some of them would be worse than none.

The original is left in storage. It is the only copy of what came out of the
PDF, it costs a few megabytes in total, and a conversion nobody can undo is not
one to run against a live bank on a Friday.

Idempotent, and a dry run by default:

    docker compose exec backend python -m scripts.convert_jpeg2000_figures
    docker compose exec backend python -m scripts.convert_jpeg2000_figures --apply
"""
import io
import sys

from sqlalchemy import text as sa_text

from app.database import SessionLocal
from app.services import storage_service, thumbnails

#: What we are converting away from, by magic bytes rather than by name — the
#: name is what got this wrong in the first place.
SUFFIXES = (".jpx", ".jp2", ".jpf")

QUALITY = 90   # Higher than a thumbnail: this replaces the figure, not a preview.


def converted_key(key: str) -> str:
    for suffix in SUFFIXES:
        if key.lower().endswith(suffix):
            return key[: -len(suffix)] + ".webp"
    return key + ".webp"


def to_webp(data: bytes) -> bytes | None:
    from PIL import Image, ImageOps

    try:
        image = ImageOps.exif_transpose(Image.open(io.BytesIO(data)))
        image = image.convert("RGBA" if image.mode in ("RGBA", "LA", "P") else "RGB")
        out = io.BytesIO()
        image.save(out, format="WEBP", quality=QUALITY, method=4)
        return out.getvalue()
    except Exception as error:
        print(f"    ! could not decode: {error}")
        return None


def rows_with_jpeg2000(db):
    """Every path that names a JPEG 2000 file, wherever it is recorded."""
    like = " OR ".join(f"lower({{col}}) LIKE '%{s}'" for s in SUFFIXES)
    found: dict[str, list[tuple[str, int, str]]] = {}
    for table, columns, key in (("questions", ("image_path", "explanation_image_path"), "id"),
                                ("media_assets", ("path",), "id")):
        for column in columns:
            clause = like.format(col=column)
            for row in db.execute(sa_text(
                    f"SELECT {key} AS id, {column} AS path FROM {table} "
                    f"WHERE {column} IS NOT NULL AND ({clause})")):
                found.setdefault(row.path, []).append((table, row.id, column))
    return found


def main(apply: bool) -> int:
    db = SessionLocal()
    try:
        found = rows_with_jpeg2000(db)
        if not found:
            print("Nothing left in JPEG 2000.")
            return 0

        print(f"{len(found)} file(s) referenced by {sum(len(v) for v in found.values())} row(s)\n")
        converted = skipped = 0
        for key, references in sorted(found.items()):
            target = converted_key(key)
            where = ", ".join(f"{t}#{i}.{c}" for t, i, c in references)
            print(f"  {key}\n    -> {target}   ({where})")

            data = storage_service.load(key)
            if not data:
                print("    ! not in storage; left alone")
                skipped += 1
                continue
            if not apply:
                continue

            if not storage_service.exists(target):
                webp = to_webp(data)
                if not webp:
                    skipped += 1
                    continue
                storage_service.save(target, webp, "image/webp")
                print(f"    {len(data):,} bytes -> {len(webp):,}")
            for table, row_id, column in references:
                db.execute(sa_text(f"UPDATE {table} SET {column} = :new WHERE {key_of(table)} = :id"),
                           {"new": target, "id": row_id})
            # The derivatives were made from the old key and are now orphaned.
            thumbnails.forget(key)
            converted += 1

        if not apply:
            print("\ndry run. Pass --apply to write.")
            return 0
        db.commit()
        print(f"\ndone. {converted} converted, {skipped} left alone. "
              f"Originals are still in storage.")
        return 0
    finally:
        db.close()


def key_of(table: str) -> str:
    return "id"


if __name__ == "__main__":
    sys.exit(main("--apply" in sys.argv))