typstcard/format.py

#!/usr/bin/env python3
import sys

sys.dont_write_bytecode = True

import argparse
import base64
import csv
import hashlib
import imb
import json
import os
import requests
import string
import subprocess
import typing
import urllib.parse
import xml.etree.ElementTree as ET

from pathlib import Path
from typing import NewType, TypedDict

_cache: Path | None = None


def cache_dir() -> Path:
    global _cache
    if _cache is None:
        _cache = root_dir() / "cache"
        _cache.mkdir(exist_ok=True)
    return _cache


_root: Path | None = None


def root_dir() -> Path:
    global _root
    if _root is None:
        _root = Path(os.path.realpath(__file__)).parent
        _root.mkdir(exist_ok=True)
    return _root


IsoCode = NewType("IsoCode", str)


def iso_code(s: str) -> IsoCode:
    if len(s) != 2:
        raise ValueError("must be 2 characters long")
    s = s.lower()
    if not (s[0] in string.ascii_lowercase and s[1] in string.ascii_lowercase):
        raise ValueError("must be ascii letters")
    return IsoCode(s)


def get_discord_avatar(
    url: urllib.parse.ParseResult, secrets: dict[str, str]
) -> typing.Optional[str]:
    try:
        uid = url.path
        token = secrets["discord_token"]
        user_info = requests.get(
            f"https://discord.com/api/users/{uid}",
            headers={"Authorization": f"Bot {token}"},
        ).json()
        avatar_hash = user_info["avatar"]
        return f"https://cdn.discordapp.com/avatars/{uid}/{avatar_hash}.png?size=4096"
    except KeyError:
        return None


def get_fedi_avatar(
    url: urllib.parse.ParseResult, secrets: dict[str, str]
) -> typing.Optional[str]:
    try:
        mastodon_api = secrets["mastodon_api"]
        user_info = requests.get(
            f"{mastodon_api}/api/v1/accounts/lookup", params={"acct": url.path}
        ).json()
        avatar_url = user_info["avatar_static"]
        return avatar_url
    except KeyError:
        return None


def get_orig_avatar(
    url: str, basename: str, secrets: dict[str, str]
) -> typing.Optional[bytes]:
    url_parts = urllib.parse.urlparse(url)
    if url_parts.scheme == "fedi":
        real_url = get_fedi_avatar(url_parts, secrets)
    elif url_parts.scheme == "discord":
        real_url = get_discord_avatar(url_parts, secrets)
    else:
        real_url = url

    if real_url is None:
        return None

    img_file = cache_dir() / basename

    if img_file.exists():
        with img_file.open("rb") as infile:
            return infile.read()
    result = requests.get(real_url)
    if not result.ok:
        return None
    with img_file.open("wb") as outfile:
        outfile.write(result.content)
    return result.content


def get_avatar(url: str, secrets: dict[str, str]) -> Path | None:
    basename = hashlib.sha256(url.encode("utf-8")).hexdigest()
    file_path = cache_dir() / f"{basename}.svg"
    if not file_path.exists():
        avatar_raster = get_orig_avatar(url, basename, secrets)
        if avatar_raster is None:
            return None

        svg_text = f"""<svg viewBox="0 0 512 512" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink">
            <clipPath id="circle">
                <circle cx="256" cy="256" r="256" />
            </clipPath>
            <image width="512" height="512" clip-path="url(#circle)"
                xlink:href="data:;base64,{base64.b64encode(avatar_raster).decode("utf-8")}" />
        </svg>"""

        with open(file_path, "w") as svgfile:
            svgfile.write(svg_text)
    return file_path.relative_to(root_dir())


def get_country_name(
    root: ET.ElementTree, destination: str, alt: str | None = None
) -> typing.Optional[str]:
    elements = root.findall(
        f"./localeDisplayNames/territories/territory[@type='{destination.upper()}']"
    )
    normal = None
    for element in elements:
        if element.attrib.get("alt") == alt:
            return element.text
        elif element.attrib.get("alt") is None:
            normal = element.text
    return normal


parser = argparse.ArgumentParser(
    prog="format", description="format postcards with latex"
)
parser.add_argument(
    "template", help="template to use", nargs="?", type=Path, default="2card"
)
parser.add_argument(
    "-o", "--origin", help="origin country code", default="us", type=iso_code
)
parser.add_argument(
    "-l",
    "--language",
    help="language to use for countries ",
    default="en",
    type=iso_code,
)
parser.add_argument(
    "-c",
    "--count",
    default=1,
    type=int,
    help="Number of sets of labels to print, default 1 (labels only)",
)
parser.add_argument(
    "-s",
    "--skip",
    default=0,
    type=int,
    help="Number of labels to skip (label sheets only)",
)

parser.add_argument(
    "-a",
    "--address-file",
    default="addresses.csv",
    type=Path,
    help="CSV file containing addresses",
)

parser.add_argument(
    "-i",
    "--content-path",
    default="content",
    type=Path,
    help="Directory containing content files",
)

parser.add_argument(
    "-n",
    "--no-content",
    action="store_true",
    help="Skip content, e.g. to make postcard back labels",
)

parser.add_argument(
    "--sheet-name",
    default="",
    type=str,
    help="Which sheet to use out of a .ods",
)

parser.add_argument(
    "-d", "--dont-compile", action="store_true", help="Don't compile to output.pdf"
)

parser.add_argument(
    "-N",
    "--no-nixowos",
    action="store_true",
    help="Don't print nixowos (e.g. if it would go over stamp)",
)

parser.add_argument(
    "-g", "--svg", action="store_true", help="Output SVG files instead of PDF"
)

parser.add_argument("-w", "--watch", action="store_true", help="Watch input files")


class Args(argparse.Namespace):
    template: Path
    origin: IsoCode
    language: IsoCode
    count: int
    skip: int
    address_file: Path
    content_path: Path
    no_content: bool
    dont_compile: bool
    sheet_name: str


args = parser.parse_args(args=None, namespace=Args())

cldr_root = ET.parse(
    f"{os.getenv('CLDR_ROOT')}/share/unicode/cldr/common/main/{args.language}.xml"
)

if args.address_file.suffix == ".csv":
    csv_filename = args.address_file
elif args.address_file.suffix == ".ods":
    # https://help.libreoffice.org/latest/en-US/text/shared/guide/csv_params.html
    export_options = (
        "csv:Text - txt - csv (StarCalc):"
        + ",".join(  # Magic CSV export string
            [
                str(ord(",")),  # Field Separator
                str(ord('"')),  # Text Delimiter
                "76",  # Character Set - UTF-8 => 76
                "",  # starting line; ignored in export
                "",  # cell format codes; ignored in export
                "1033",  # Language Identifier - en-US => 1033
                "",  # Quoted field as text; default False
                "",  # Detect special numbers; default True
                "",  # Save cell contents as shown; default True
                "",  # Export cell formulas; default false
                "",  # Remove spaces; ignored in export
                "-1",  # Export sheets - all sheets => -1
            ]
        )
    )
    result = subprocess.run(
        [
            "libreoffice",
            "--headless",
            "--convert-to",
            export_options,
            "--outdir",
            str(cache_dir()),
            args.address_file,
        ]
    )
    assert result.returncode == 0
    csv_filename = cache_dir() / f"{args.address_file.stem}-{args.sheet_name}.csv"
else:
    raise Exception("Unknown file type for --address-file")


class CSVRow(TypedDict):
    Name: str
    Address: str
    DPC: str
    Country: str
    Design: str
    Type: str
    Personalization: str
    Avatar: str | None
    US: str


csvfile = open(csv_filename)
rows = csv.DictReader(csvfile)

with open("secrets.json") as secrets_file:
    secrets = json.load(secrets_file)

current_serial = imb.get_first_serial()
mid = secrets.get("mailer_id")


class Card(TypedDict):
    address: str
    avatar: Path | None
    row: CSVRow
    imb: str


cards: list[Card] = []
for row in rows:
    row = CSVRow(**row)
    if row["Address"] == "":
        continue

    if row["Country"].lower() == args.origin:
        country = []
    else:
        name = get_country_name(cldr_root, row["Country"])
        assert name is not None
        country = [name.upper()]

    address = row["Address"].split("\n") + country

    if (avatar_url := row.get("Avatar", "")) != "":
        avatar = get_avatar(avatar_url, secrets)  # type: ignore
    else:
        avatar = None

    cards.append(
        {
            "address": "\n".join(address),
            "avatar": avatar,
            "row": row,
            "imb": "",
        }
    )

# Typst can't access files outside the project root, except through a symlink
# Create one in cache to use here
p = cache_dir() / "content"
p.unlink(missing_ok=True)
content_full_path = Path(os.getcwd()).joinpath(args.content_path)
p.symlink_to(content_full_path)

cards = cards * args.count

serial = imb.get_first_serial()
if args.origin == "us":
    for card in cards:
        dpc = card["row"].get("DPC", "")
        if dpc != "" and mid is not None:
            card["imb"] = imb.generate(
                0, 310, mid, serial, dpc.replace(" ", "").replace("-", "")
            )
            serial += 1
imb.write_current_serial(serial)


def serialize_paths(obj: object):
    if isinstance(obj, Path):
        return str(obj)
    raise TypeError("Type not Serializable")


with (cache_dir() / "options.json").open("w") as options:
    json.dump(
        fp=options,
        obj={
            "args": args.__dict__,
            "cards": cards,
        },
        default=serialize_paths,
    )

if args.dont_compile:
    exit()

font_paths = os.getenv("TYPST_FONT_PATHS")
assert font_paths is not None

if args.svg:
    if os.path.exists(root_dir() / "output"):
        for f in (root_dir() / "output").glob("*.svg"):
            os.remove(f)
    else:
        os.mkdir(root_dir() / "output")

typst_args = [
    "typst",
    "watch" if args.watch else "compile",
    "--root",
    root_dir(),
    "--font-path",
    cache_dir() / "content",
    "--font-path",
    font_paths,
    args.template,
    root_dir() / "output" / "{n}.svg" if args.svg else "output.pdf",
]

# print(*typst_args)

os.execlp(
    "typst",
    *typst_args,
)