#!/usr/bin/env python3
"""Generate a QRBackup-style paper backup PDF for a file.

The QR payload format is intentionally simple and printable:

    QB1/<index>/<total>/GZ/<raw_size>/<sha256>/<base45_chunk>/END

Each payload only uses QR alphanumeric characters so QR codes stay compact.
The generated restore script can rebuild the original file from scanned payloads.
"""

from __future__ import annotations

import argparse
import gzip
import hashlib
import json
import math
import textwrap
from pathlib import Path

from PIL import Image, ImageDraw, ImageFont
import qrcode
from qrcode.exceptions import DataOverflowError


BASE45_ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ $%*+-./:"


def base45_encode(data: bytes) -> str:
    out: list[str] = []
    i = 0
    while i + 1 < len(data):
        value = data[i] * 256 + data[i + 1]
        out.append(BASE45_ALPHABET[value % 45])
        out.append(BASE45_ALPHABET[(value // 45) % 45])
        out.append(BASE45_ALPHABET[value // (45 * 45)])
        i += 2
    if i < len(data):
        value = data[i]
        out.append(BASE45_ALPHABET[value % 45])
        out.append(BASE45_ALPHABET[value // 45])
    return "".join(out)


def chunk_text(text: str, size: int) -> list[str]:
    return [text[i : i + size] for i in range(0, len(text), size)]


def make_payloads(encoded: str, chunk_size: int, raw_size: int, sha256: str, compression: str) -> list[str]:
    chunks = chunk_text(encoded, chunk_size)
    total = len(chunks)
    return [
        f"QB1/{index:04d}/{total:04d}/{compression}/{raw_size}/{sha256}/{chunk}/END"
        for index, chunk in enumerate(chunks, start=1)
    ]


ERROR_CORRECTION_LEVELS = {
    "L": qrcode.constants.ERROR_CORRECT_L,
    "M": qrcode.constants.ERROR_CORRECT_M,
    "Q": qrcode.constants.ERROR_CORRECT_Q,
    "H": qrcode.constants.ERROR_CORRECT_H,
}


def build_qr(payload: str, box_size: int = 5, error_correction: str = "Q") -> Image.Image:
    qr = qrcode.QRCode(
        version=40,
        error_correction=ERROR_CORRECTION_LEVELS[error_correction],
        box_size=box_size,
        border=4,
    )
    qr.add_data(payload, optimize=0)
    qr.make(fit=False)
    return qr.make_image(fill_color="black", back_color="white").convert("RGB")


def payloads_that_fit(encoded: str, raw_size: int, sha256: str, start_chunk_size: int, error_correction: str, compression: str) -> tuple[int, list[str]]:
    chunk_size = start_chunk_size
    while chunk_size >= 800:
        payloads = make_payloads(encoded, chunk_size, raw_size, sha256, compression)
        try:
            for payload in payloads:
                build_qr(payload, box_size=1, error_correction=error_correction)
            return chunk_size, payloads
        except DataOverflowError:
            chunk_size -= 100
    raise RuntimeError("Unable to fit payloads into QR version 40-Q codes")


def load_font(size: int) -> ImageFont.ImageFont:
    candidates = [
        "arial.ttf",
        "calibri.ttf",
        "C:/Windows/Fonts/arial.ttf",
        "C:/Windows/Fonts/calibri.ttf",
    ]
    for candidate in candidates:
        try:
            return ImageFont.truetype(candidate, size=size)
        except Exception:
            pass
    return ImageFont.load_default()


def draw_wrapped(draw: ImageDraw.ImageDraw, xy: tuple[int, int], text: str, font: ImageFont.ImageFont, width: int, fill=(0, 0, 0)) -> int:
    x, y = xy
    for paragraph in text.splitlines():
        if not paragraph:
            y += 24
            continue
        for line in textwrap.wrap(paragraph, width=width):
            draw.text((x, y), line, font=font, fill=fill)
            y += 30
    return y


def make_pages(payloads: list[str], source: Path, raw_size: int, payload_size: int, sha256: str, per_page: int, error_correction: str, compression: str) -> list[Image.Image]:
    page_w, page_h = 2480, 3508  # A4 at 300 DPI
    margin = 115
    header_h = 260
    gap_x = 90
    gap_y = 95
    cols = 2
    rows = max(1, math.ceil(per_page / cols))
    cell_w = (page_w - 2 * margin - gap_x) // cols
    cell_h = (page_h - header_h - 2 * margin - (rows - 1) * gap_y) // rows

    title_font = load_font(46)
    body_font = load_font(24)
    small_font = load_font(22)
    label_font = load_font(28)
    pages: list[Image.Image] = []
    total_pages = math.ceil(len(payloads) / per_page)

    for page_index in range(total_pages):
        page = Image.new("RGB", (page_w, page_h), "white")
        draw = ImageDraw.Draw(page)
        draw.text((margin, 70), "QRBACKUP PAPER BACKUP", font=title_font, fill=(0, 0, 0))
        info = (
            f"FILE: {source.name}    QR: {len(payloads)}    PAGE: {page_index + 1}/{total_pages}\n"
            f"RAW: {raw_size} BYTES    PAYLOAD: {payload_size} BYTES    MODE: {compression}    ECC: {error_correction}    SHA256: {sha256}"
        )
        draw_wrapped(draw, (margin, 135), info, body_font, width=138)

        page_payloads = payloads[page_index * per_page : (page_index + 1) * per_page]
        for local_index, payload in enumerate(page_payloads):
            row = local_index // cols
            col = local_index % cols
            x0 = margin + col * (cell_w + gap_x)
            y0 = header_h + margin + row * (cell_h + gap_y)
            global_index = page_index * per_page + local_index + 1

            qr_img = build_qr(payload, box_size=5, error_correction=error_correction)
            qr_x = x0 + (cell_w - qr_img.width) // 2
            qr_y = y0 + 12
            page.paste(qr_img, (qr_x, qr_y))

            label = f"QR {global_index:04d}/{len(payloads):04d}"
            label_bbox = draw.textbbox((0, 0), label, font=label_font)
            label_w = label_bbox[2] - label_bbox[0]
            draw.text((x0 + (cell_w - label_w) // 2, qr_y + qr_img.height + 18), label, font=label_font, fill=(0, 0, 0))

            short_hash = hashlib.sha256(payload.encode("ascii")).hexdigest().upper()[:16]
            draw.text((x0 + 12, y0 + cell_h - 34), f"PAYLOAD-SHA256-16: {short_hash}", font=small_font, fill=(0, 0, 0))

        footer = "Scan all QR payloads, one per line, then run restore_from_qr_payloads.py."
        draw.text((margin, page_h - 72), footer, font=body_font, fill=(0, 0, 0))
        pages.append(page)

    return pages


def write_restore_script(path: Path) -> None:
    script = r'''#!/usr/bin/env python3
import argparse
import gzip
import hashlib
from pathlib import Path

BASE45_ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ $%*+-./:"
BASE45_INDEX = {ch: i for i, ch in enumerate(BASE45_ALPHABET)}

def base45_decode(text: str) -> bytes:
    out = bytearray()
    i = 0
    while i < len(text):
        if i + 2 < len(text):
            value = BASE45_INDEX[text[i]] + BASE45_INDEX[text[i + 1]] * 45 + BASE45_INDEX[text[i + 2]] * 45 * 45
            out.extend(divmod(value, 256))
            i += 3
        elif i + 1 < len(text):
            value = BASE45_INDEX[text[i]] + BASE45_INDEX[text[i + 1]] * 45
            out.append(value)
            i += 2
        else:
            raise ValueError("Invalid base45 length")
    return bytes(out)

def parse_payload(line: str):
    tag, index, total, compression, raw_size, sha256, chunk_and_end = line.strip().split("/", 6)
    if tag != "QB1":
        raise ValueError(f"Unsupported payload tag: {tag}")
    if not chunk_and_end.endswith("/END"):
        raise ValueError("Payload missing /END sentinel")
    chunk = chunk_and_end[:-4]
    return int(index), int(total), compression, int(raw_size), sha256.upper(), chunk

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("payload_text", help="Text file containing scanned QR payloads, one per line")
    parser.add_argument("-o", "--output", default="PROJECT_REPLICATION_PROMPT.restored.md")
    args = parser.parse_args()

    entries = []
    for raw_line in Path(args.payload_text).read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line:
            continue
        if not line.startswith("QB1/"):
            continue
        entries.append(parse_payload(line))

    if not entries:
        raise SystemExit("No QB1 payloads found")

    total = entries[0][1]
    compression = entries[0][2]
    raw_size = entries[0][3]
    expected_sha = entries[0][4]
    by_index = {}
    for index, item_total, item_compression, item_raw_size, item_sha, chunk in entries:
        if item_total != total or item_compression != compression or item_raw_size != raw_size or item_sha != expected_sha:
            raise SystemExit("Payload metadata mismatch")
        by_index[index] = chunk

    missing = [i for i in range(1, total + 1) if i not in by_index]
    if missing:
        raise SystemExit(f"Missing QR payloads: {missing}")

    encoded = "".join(by_index[i] for i in range(1, total + 1))
    payload_data = base45_decode(encoded)
    if compression == "GZ":
        data = gzip.decompress(payload_data)
    elif compression == "RAW":
        data = payload_data
    else:
        raise SystemExit(f"Unsupported compression: {compression}")
    actual_sha = hashlib.sha256(data).hexdigest().upper()
    if len(data) != raw_size or actual_sha != expected_sha:
        raise SystemExit(f"Verification failed: size={len(data)} sha256={actual_sha}")
    Path(args.output).write_bytes(data)
    print(f"Restored {args.output} ({len(data)} bytes, sha256={actual_sha})")

if __name__ == "__main__":
    main()
'''
    path.write_text(script, encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("input", type=Path)
    parser.add_argument("-o", "--output-dir", type=Path, default=Path("qrbackup_output"))
    parser.add_argument("--chunk-size", type=int, default=2200)
    parser.add_argument("--per-page", type=int, default=4)
    parser.add_argument("--ecc", choices=sorted(ERROR_CORRECTION_LEVELS), default="Q", help="QR error correction level: L has the fewest codes, H is most robust")
    parser.add_argument("--no-gzip", action="store_true", help="Store the input bytes directly; useful when the input is already compressed")
    args = parser.parse_args()

    source = args.input.resolve()
    output_dir = args.output_dir.resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    raw = source.read_bytes()
    compressed = raw if args.no_gzip else gzip.compress(raw, compresslevel=9)
    compression = "RAW" if args.no_gzip else "GZ"
    sha256 = hashlib.sha256(raw).hexdigest().upper()
    encoded = base45_encode(compressed)
    chunk_size, payloads = payloads_that_fit(encoded, len(raw), sha256, args.chunk_size, args.ecc, compression)

    base_name = source.stem
    payload_text = output_dir / f"{base_name}.qrpayloads.txt"
    payload_text.write_text("\n".join(payloads) + "\n", encoding="ascii")

    manifest = {
        "format": "QB1",
        "source": str(source),
        "file_name": source.name,
        "raw_size": len(raw),
        "payload_size": len(compressed),
        "compression": compression,
        "base45_size": len(encoded),
        "chunk_size": chunk_size,
        "qr_count": len(payloads),
        "per_page": args.per_page,
        "sha256": sha256,
        "error_correction": args.ecc,
        "qr_version": 40,
    }
    manifest_path = output_dir / f"{base_name}.qrbackup_manifest.json"
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")

    restore_script = output_dir / "restore_from_qr_payloads.py"
    write_restore_script(restore_script)

    readme = output_dir / "README_RESTORE.txt"
    readme.write_text(
        "\n".join(
            [
                "QRBACKUP PAPER BACKUP",
                "",
                f"Original file: {source.name}",
                f"Original size: {len(raw)} bytes",
                f"Gzip size: {len(compressed)} bytes",
                f"QR count: {len(payloads)}",
                f"SHA256: {sha256}",
                "",
                "Restore steps:",
                "1. Scan every QR code and save each QR payload as one line in a text file.",
                "2. Run: python restore_from_qr_payloads.py scanned_payloads.txt -o PROJECT_REPLICATION_PROMPT.restored.md",
                "3. The script verifies size and SHA256 before writing the restored file.",
                "",
                "The generated .qrpayloads.txt file contains the exact QR payloads and can be used to self-test the restore script.",
            ]
        ),
        encoding="utf-8",
    )

    pages = make_pages(payloads, source, len(raw), len(compressed), sha256, args.per_page, args.ecc, compression)
    combined_pdf = output_dir / f"{base_name}.qrbackup.all_pages.pdf"
    pages[0].save(combined_pdf, "PDF", save_all=True, append_images=pages[1:], resolution=300)

    for index, page in enumerate(pages, start=1):
        page_pdf = output_dir / f"{base_name}.qrbackup.page_{index:03d}.pdf"
        page.save(page_pdf, "PDF", resolution=300)

    print(json.dumps({**manifest, "output_dir": str(output_dir), "combined_pdf": str(combined_pdf), "pages": len(pages)}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()
