#!/usr/bin/env python3
"""Check a comic archive before you translate it.

Usage:
    python3 check_comic_archive.py BOOK.cbz
    python3 check_comic_archive.py BOOK.cbz --repack FIXED.cbz
    python3 check_comic_archive.py BOOK.cbz --first cover.jpg --repack FIXED.cbz
    python3 check_comic_archive.py EXTRACTED_FOLDER --repack BOOK.cbz

What it reports:
  * the container type (CBZ is a ZIP; CBR, CB7 and CBT are RAR, 7z and tar);
  * the page order a numeric-aware ("natural") sort gives, and where a plain
    character-by-character sort disagrees with it;
  * pages that will land out of place (an unnumbered cover.jpg sorts after
    001.jpg), unpadded numbers, junk entries, images that are not a page
    format, two-page spreads and very tall strips;
  * the ComicInfo.xml reading direction and page table, if present.

--repack writes a new CBZ whose page images are renamed 0001.ext, 0002.ext, ...
in the natural order shown (after moving any --first pages to the front),
with junk entries dropped and every other file
(ComicInfo.xml included) copied unchanged. Check the printed order first:
renaming bakes it in.

Python 3.8+ standard library only. Nothing is uploaded anywhere.
"""

import argparse
import os
import re
import struct
import sys
import xml.etree.ElementTree as ET
import zipfile

PAGE_EXTS = {"png", "jpg", "jpeg", "webp", "gif", "bmp", "tif", "tiff"}
OTHER_IMAGE_EXTS = {"avif", "jxl", "heic", "heif", "svg", "jp2"}


def ext_of(name):
    m = re.search(r"\.([A-Za-z0-9]+)$", name)
    return m.group(1).lower() if m else ""


def is_junk(name):
    return any(p.startswith(".") or p == "__MACOSX" or p.lower() == "thumbs.db"
               for p in name.split("/"))


# Approximates Intl.Collator("en", {numeric: true, sensitivity: "base"}) for
# ASCII names: spaces and punctuation sort before digits, digits before
# letters, digit runs compare by value and letters ignore case. Accented and
# non-Latin names can sort differently; zero-padded ASCII names avoid that.
PUNCTUATION = " _-,;:!?.'\"()[]{}@*/\\&#%`^+<=>|~$"


def natural_key(name):
    key = []
    for run in re.findall(r"\d+|.", name):
        if run.isdigit():
            key.append((1, int(run)))
        elif run.isalpha():
            key.append((2, run.lower()))
        else:
            i = PUNCTUATION.find(run)
            key.append((0, i if i >= 0 else len(PUNCTUATION) + ord(run)))
    return key


def sniff_container(head):
    if head.startswith(b"PK\x03\x04") or head.startswith(b"PK\x05\x06"):
        return "zip"
    if head.startswith(b"Rar!\x1a\x07"):
        return "rar"
    if head.startswith(b"7z\xbc\xaf\x27\x1c"):
        return "7z"
    if len(head) > 262 and head[257:262] == b"ustar":
        return "tar"
    return "unknown"


def image_size(data):
    """Width and height from the file header, or None."""
    if data[:8] == b"\x89PNG\r\n\x1a\n" and len(data) >= 24:
        return struct.unpack(">II", data[16:24])
    if data[:6] in (b"GIF87a", b"GIF89a"):
        return struct.unpack("<HH", data[6:10])
    if data[:2] == b"BM" and len(data) >= 26:
        w, h = struct.unpack("<ii", data[18:26])
        return w, abs(h)
    if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        kind = data[12:16]
        if kind == b"VP8X":
            w = int.from_bytes(data[24:27], "little") + 1
            h = int.from_bytes(data[27:30], "little") + 1
            return w, h
        if kind == b"VP8 ":
            w, h = struct.unpack("<HH", data[26:30])
            return w & 0x3FFF, h & 0x3FFF
        if kind == b"VP8L":
            b = data[21:25]
            w = 1 + (((b[1] & 0x3F) << 8) | b[0])
            h = 1 + (((b[3] & 0x0F) << 10) | (b[2] << 2) | ((b[1] & 0xC0) >> 6))
            return w, h
    if data[:2] == b"\xff\xd8":
        i = 2
        while i + 9 < len(data):
            if data[i] != 0xFF:
                i += 1
                continue
            marker = data[i + 1]
            if marker in (0xC0, 0xC1, 0xC2, 0xC3, 0xC5, 0xC6, 0xC7,
                          0xC9, 0xCA, 0xCB, 0xCD, 0xCE, 0xCF):
                h, w = struct.unpack(">HH", data[i + 5:i + 9])
                return w, h
            if marker in (0xD8, 0x01) or 0xD0 <= marker <= 0xD7:
                i += 2
                continue
            i += 2 + struct.unpack(">H", data[i + 2:i + 4])[0]
    return None


def read_entries(path):
    """Return {name: bytes} for a CBZ file or an extracted folder."""
    if os.path.isdir(path):
        out = {}
        for root, _, files in os.walk(path):
            for f in files:
                full = os.path.join(root, f)
                rel = os.path.relpath(full, path).replace(os.sep, "/")
                with open(full, "rb") as fh:
                    out[rel] = fh.read()
        return out
    with open(path, "rb") as fh:
        head = fh.read(300)
    kind = sniff_container(head)
    if kind != "zip":
        hint = {
            "rar": "This is a RAR archive (CBR). Extract it with 7-Zip, unar or "
                   "The Unarchiver, then run this script on the folder with "
                   "--repack to get a CBZ.",
            "7z": "This is a 7z archive (CB7). Extract it first, then run this "
                  "script on the folder with --repack.",
            "tar": "This is a tar archive (CBT). Extract it first, then run this "
                   "script on the folder with --repack.",
        }.get(kind, "This file is not a ZIP, RAR, 7z or tar archive.")
        sys.exit(f"Container: {kind}\n{hint}")
    try:
        with zipfile.ZipFile(path) as z:
            return {i.filename: z.read(i) for i in z.infolist() if not i.is_dir()}
    except zipfile.BadZipFile as e:
        sys.exit(f"Container: zip, but it could not be read ({e}). The file may be "
                 "incomplete or damaged; download or export it again.")


def comicinfo_report(data, page_count):
    lines = []
    try:
        root = ET.fromstring(data)
    except ET.ParseError as e:
        return [f"  ComicInfo.xml is not valid XML ({e})."]
    manga = (root.findtext("Manga") or "").strip()
    lines.append(f"  Manga (reading direction): {manga or 'not set'}")
    if manga == "YesAndRightToLeft":
        lines.append("  Right-to-left is stored here, not in the images. Keep this "
                     "file in the archive so readers that honor it still see it.")
    lang = (root.findtext("LanguageISO") or "").strip()
    if lang:
        lines.append(f"  LanguageISO: {lang}. Update it after translation; a copied "
                     "ComicInfo.xml still names the source language.")
    pages = root.find("Pages")
    if pages is not None:
        rows = list(pages)
        lines.append(f"  Page table: {len(rows)} entries for {page_count} page images.")
        if len(rows) != page_count:
            lines.append("  ! The page table and the image count disagree.")
        sized = [r for r in rows if r.get("ImageSize") or r.get("ImageWidth")]
        if sized:
            lines.append(f"  {len(sized)} entries record ImageSize/ImageWidth. These "
                         "describe the original files and no longer match once "
                         "pages are replaced.")
    return lines


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("source", help="a .cbz file or an extracted folder")
    ap.add_argument("--repack", metavar="OUT.cbz",
                    help="write a CBZ with zero-padded page names in natural order")
    ap.add_argument("--first", metavar="NAME", action="append", default=[],
                    help="move this page to the front before repacking (repeatable)")
    args = ap.parse_args()

    entries = read_entries(args.source)
    junk = sorted(n for n in entries if is_junk(n))
    live = {n: b for n, b in entries.items() if not is_junk(n)}
    pages = [n for n in live if ext_of(n) in PAGE_EXTS]
    other_images = sorted(n for n in live if ext_of(n) in OTHER_IMAGE_EXTS)
    others = sorted(n for n in live if n not in pages and n not in other_images)

    natural = sorted(pages, key=natural_key)
    plain = sorted(pages)
    missing = [n for n in args.first if n not in pages]
    if missing:
        sys.exit("--first: no page named " + ", ".join(missing))

    print(f"Source: {args.source}")
    print(f"Page images: {len(pages)}   other files: {len(others)}   "
          f"junk entries: {len(junk)}")
    print()
    print("Reading order (natural sort):")
    for i, n in enumerate(natural, 1):
        size = image_size(live[n])
        note = ""
        if size:
            w, h = size
            note = f"{w}x{h}"
            if w > h:
                note += "  spread (wider than tall)"
            elif h > 3 * w:
                note += "  tall strip (more than 3x height)"
        print(f"  {i:>4}  {n}  {note}")

    problems = []
    if natural != plain:
        first = next(i for i, (a, b) in enumerate(zip(natural, plain)) if a != b)
        problems.append(
            f"A plain character sort puts '{plain[first]}' at position {first + 1}, "
            f"where natural sort has '{natural[first]}'. A reader without "
            "natural sorting will show pages out of order.")
    numbered = [n for n in natural if re.search(r"\d", os.path.basename(n))]
    unnumbered = [n for n in natural if n not in numbered]
    if numbered and unnumbered:
        where = ", ".join(f"'{n}' at {natural.index(n) + 1}" for n in unnumbered)
        problems.append(
            f"Pages without a number in the file name: {where}. Digits sort "
            "before letters, so an unnumbered cover usually ends up last.")
    widths = {len(m) for n in numbered
              for m in re.findall(r"\d+", os.path.basename(n))[-1:]}
    if len(widths) > 1:
        problems.append(
            f"Page numbers use {sorted(widths)} digits. Zero-pad them "
            "(001, 002, ... 010) so every reader agrees on the order.")
    stems = {}
    for n in pages:
        stems.setdefault(re.sub(r"\.[^./]+$", "", n), []).append(n)
    clashes = [v for v in stems.values() if len(v) > 1]
    if clashes:
        problems.append(
            "Same name, different extension: "
            + "; ".join(" + ".join(sorted(v)) for v in clashes)
            + ". Tools that re-encode pages to one format must rename one of them.")
    if other_images:
        problems.append(
            "Images in a format that is not treated as a page here: "
            + ", ".join(other_images)
            + ". Convert them to PNG, JPEG or WebP, or they stay untranslated.")
    folders = sorted({n.rsplit("/", 1)[0] for n in pages if "/" in n})
    if len(folders) > 1:
        problems.append(
            f"Pages sit in {len(folders)} folders; the order above sorts the full "
            "path, so check that chapter folders are numbered the same way.")

    print()
    if junk:
        print("Junk entries (skipped): " + ", ".join(junk))
    if others:
        print("Other files (carried over unchanged): " + ", ".join(others))
    info = next((n for n in live if n.lower().endswith("comicinfo.xml")), None)
    if info:
        print("ComicInfo.xml:")
        print("\n".join(comicinfo_report(live[info], len(pages))))
    print()
    if problems:
        print("Check before uploading:")
        for p in problems:
            print("  - " + p)
    else:
        print("No ordering or content problems found.")

    if args.repack:
        order = args.first + [n for n in natural if n not in args.first]
        width = max(4, len(str(len(order))))
        with zipfile.ZipFile(args.repack, "w", zipfile.ZIP_STORED) as out:
            for i, n in enumerate(order, 1):
                out.writestr(f"{i:0{width}d}.{ext_of(n)}", live[n])
            for n in other_images + others:
                # Pages are flattened to the root, so keep ComicInfo.xml there too.
                name = "ComicInfo.xml" if n == info else n
                out.writestr(name, live[n])
        print()
        if args.first and info:
            print("Note: --first changed page positions. Check that the <Pages> "
                  "entries in ComicInfo.xml (for example Type=\"FrontCover\") "
                  "still point to the right Image numbers.")
        print(f"Wrote {args.repack}: {len(order)} pages renamed "
              f"{1:0{width}d}..{len(order):0{width}d}, "
              f"{len(junk)} junk entries dropped.")


if __name__ == "__main__":
    main()
