typstcard/format.py

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

from pathlib import Path
from typing import NewType, TypedDict

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


_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]) -> str | 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 str(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="?", 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")


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


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"
)

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
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)

with (root_dir() / "options.json").open("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

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

# print(*typst_args)

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