typstcard/format.py

#!/usr/bin/env python3
import argparse
import base64
import csv
import hashlib
import json
import os
import pathlib
import string
import typing
import urllib.parse
import xml.etree.ElementTree as ET
import requests
import imb

from pathlib import Path
from typing import TypedDict

# A lot of stuff needs to be in the same directory, just chdir
os.chdir(os.path.dirname(os.path.realpath(__file__)))


def cache_dir() -> Path:
    cache = Path("cache")
    cache.mkdir(exist_ok=True)
    return cache


def iso_code(s: str) -> str:
    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 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]) -> str | None:
    basename = hashlib.sha256(url.encode("utf-8")).hexdigest()
    file_path = cache_dir() / f"{basename}.svg"
    if file_path.exists():
        return str(file_path)
    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 str(file_path)


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="?", 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=str,
    help="CSV file containing addresses",
)

parser.add_argument(
    "-i",
    "--content-path",
    default="content",
    type=str,
    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(
    "-d", "--dont-compile", action="store_true", help="Don't compile to output.pdf"
)

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

args = parser.parse_args()

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

csvfile = open(args.address_file)
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: str | None
    row: dict[str, str]
    imb: str


cards: list[Card] = []
for row in rows:
    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 row.get("Avatar", "") != "":
        avatar = get_avatar(row["Avatar"], secrets)
    else:
        avatar = None

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

# Typst can't access files outside the project root, except through a symlink
# Create one in cache to use here
if not os.path.exists("cache"):
    os.mkdir("cache")
p = pathlib.Path("cache/content")
p.unlink(missing_ok=True)
p.symlink_to(args.content_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)

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

if args.dont_compile:
    exit()

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

os.execlp(
    "typst",
    "typst",
    "watch" if args.watch else "compile",
    "--font-path",
    args.content_path,
    "--font-path",
    font_paths,
    args.template,
    "output.pdf",
)