"""Font bank over the openly licensed fonts fetched by scripts/fetch_fonts.py.

Families are split deterministically into train / val so evaluation measures unseen fonts.
"""

import zlib
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path

from fontTools.ttLib import TTFont
from PIL import ImageFont

MONO_HINTS = ("mono", "courier", "vt323", "inconsolata", "anonymouspro", "cousine", "sharetech", "dotgothic", "spacemono", "ptmono")
VAL_PERCENT = 15
# Maths / symbol fonts: never a line's text font, only the fallback for glyphs it lacks (in every split).
FALLBACK_FAMILIES = ("notosansmath", "stixtwotext")


@dataclass(frozen=True)
class FontFace:
    path: str
    family: str
    mono: bool
    italic: bool
    variable: bool
    bold_file: bool


@lru_cache(maxsize=None)
def coverage(path: str) -> frozenset:
    font = TTFont(path, lazy=True, fontNumber=0)
    cmap = font.getBestCmap() or {}
    return frozenset(chr(c) for c in cmap)


@lru_cache(maxsize=4096)
def load_font(path: str, size: int, bold: bool) -> ImageFont.FreeTypeFont:
    font = ImageFont.truetype(path, size)
    if "[" in Path(path).name:
        try:
            axes = font.get_variation_axes()
            values = []
            for ax in axes:
                name = ax.get("name", b"")
                name = name.decode() if isinstance(name, bytes) else str(name)
                if name.lower().startswith("weight"):
                    values.append(max(ax["minimum"], min(ax["maximum"], 700 if bold else 400)))
                else:
                    values.append(ax["default"])
            font.set_variation_by_axes(values)
        except (OSError, KeyError):
            pass
    return font


class FontBank:
    def __init__(self, root="fonts", split="train"):
        faces, self.fallbacks = [], []
        for fam_dir in sorted(p for p in Path(root).glob("*/*") if p.is_dir()):
            fam = fam_dir.name
            if fam in FALLBACK_FAMILIES:
                self.fallbacks += [FontFace(str(f), fam, False, False, "[" in f.name, False)
                                   for f in sorted(fam_dir.glob("*.[ot]tf")) if not f.name.startswith("._")
                                   and "Italic" not in f.name and ("Bold" not in f.name or "[" in f.name)]
                continue
            is_val = zlib.crc32(fam.encode()) % 100 < VAL_PERCENT
            if split != "all" and is_val != (split == "val"):
                continue
            for f in sorted(f for f in fam_dir.glob("*.[ot]tf") if not f.name.startswith("._")):   # skip macOS metadata files
                n = f.name
                faces.append(FontFace(
                    path=str(f), family=fam,
                    mono=any(h in fam for h in MONO_HINTS),
                    italic="Italic" in n,
                    variable="[" in n,
                    bold_file="Bold" in n and not any(k in n for k in ("Semi", "Extra", "Ultra")),
                ))
        if not faces:
            raise FileNotFoundError(f"no fonts under {root}/ for split={split}; run scripts/fetch_fonts.py")
        self.faces = faces
        self.families = sorted({f.family for f in faces})

    def pick(self, rng, mono=None, bold=False, italic_ok=False):
        """Returns (face, bold_flag). Prefers a real bold file, else a variable font's weight axis."""
        fams = [f for f in self.families if mono is None or any(h in f for h in MONO_HINTS) == mono] or self.families
        fam = fams[rng.integers(len(fams))]
        cands = [f for f in self.faces if f.family == fam and (italic_ok or not f.italic)] or \
                [f for f in self.faces if f.family == fam]
        if bold:
            bolds = [f for f in cands if f.bold_file]
            if bolds:
                return bolds[rng.integers(len(bolds))], False
            var = [f for f in cands if f.variable]
            if var:
                return var[0], True
        regular = [f for f in cands if not f.bold_file] or cands
        return regular[rng.integers(len(regular))], False

    @staticmethod
    def font(face: FontFace, size: int, bold_axis: bool = False):
        return load_font(face.path, int(size), bold_axis)

    def fallback(self, ch: str):
        """A maths / symbol face that has `ch`, for glyphs the text font lacks (None when none has it)."""
        for f in self.fallbacks:
            if ch in coverage(f.path):
                return f
        return None

    @staticmethod
    def supported(face: FontFace, text: str) -> str:
        cov = coverage(face.path)
        return "".join(c if (c in cov or c == " ") else "" for c in text)