#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = ["pillow>=10"]
# ///
"""Render fastfetch output to an image (png, jpg, webp, gif). Colors are kept."""

from __future__ import annotations

import argparse
import getpass
import os
import pty
import re
import select
import shlex
import shutil
import signal
import socket
import subprocess
import sys
import termios
import time
import tty
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import cast, final

from PIL import Image, ImageDraw, ImageFont

Color = int | tuple[int, int, int]
RGB = tuple[int, int, int]

CACHE_DIR = Path.home() / ".cache" / "fetch-snap"
ALIAS_ENV = "_FSC_ALIAS"
DEFAULT_BG: RGB = (18, 18, 18)
DEFAULT_FG: RGB = (230, 230, 230)

CAPTURE_COLS = 160
CAPTURE_QUIET = 1.5
CAPTURE_DEADLINE = 15.0
PROMPT_QUIET = 0.3
PROMPT_DEADLINE = 3.0
OSC_PROBE_TIMEOUT = 0.25
SCAN_DEPTH = 10

# ---------------------------------------------------------------- palette

STANDARD: list[RGB] = [
    (0, 0, 0),
    (128, 0, 0),
    (0, 128, 0),
    (128, 128, 0),
    (0, 0, 128),
    (128, 0, 128),
    (0, 128, 128),
    (192, 192, 192),
    (128, 128, 128),
    (255, 0, 0),
    (0, 255, 0),
    (255, 255, 0),
    (0, 0, 255),
    (255, 0, 255),
    (0, 255, 255),
    (255, 255, 255),
]


def xterm256(n: int) -> RGB:
    if n < 16:
        return STANDARD[n]
    if n <= 231:
        n -= 16
        steps = (0, 95, 135, 175, 215, 255)
        return steps[n // 36], steps[(n // 6) % 6], steps[n % 6]
    v = 8 + (n - 232) * 10
    return v, v, v


# ---------------------------------------------------------------- ansi model


@dataclass
class Attr:
    fg: Color | None = None
    bg: Color | None = None
    bold: bool = False
    inverse: bool = False


@dataclass
class Seg:
    text: str
    attr: Attr


@dataclass
class Theme:
    fg: Color | None = None
    bg: Color | None = None
    colors: list[Color | None] = field(default_factory=lambda: [None] * 16)
    extended: dict[int, Color] = field(default_factory=dict)

    @property
    def background(self) -> RGB:
        return self.bg if isinstance(self.bg, tuple) else DEFAULT_BG

    @property
    def foreground(self) -> RGB:
        return self.fg if isinstance(self.fg, tuple) else DEFAULT_FG


# ANSI escape grammar: CSI = ESC [ params/intermediates final
ANSI_TOKEN_RE = re.compile(
    r"\x1b\[[\x20-\x3f]*[\x40-\x7e]|\x1bP[^\x1b]*(?:\x1b\\)?|\x1b\][^\x1b\x07]*(?:\x07|\x1b\\)|[\s\S]"
)


def apply_sgr(a: Attr, body: str) -> None:
    if not body:
        a.fg = a.bg = None
        a.bold = a.inverse = False
        return
    parts = body.split(";")
    i = 0
    while i < len(parts):
        v = int(parts[i]) if parts[i] else None
        if v is None:
            i += 1
            continue
        if v == 0:
            a.fg = a.bg = None
            a.bold = a.inverse = False
        elif v == 1:
            a.bold = True
        elif v == 22:
            a.bold = False
        elif v == 7:
            a.inverse = True
        elif v == 27:
            a.inverse = False
        elif 30 <= v <= 37:
            a.fg = v - 30
        elif 90 <= v <= 97:
            a.fg = v - 90 + 8
        elif 40 <= v <= 47:
            a.bg = v - 40
        elif 100 <= v <= 107:
            a.bg = v - 100 + 8
        elif v == 39:
            a.fg = None
        elif v == 49:
            a.bg = None
        elif v in (38, 48) and i + 2 < len(parts) and parts[i + 1] == "5":
            if v == 38:
                a.fg = int(parts[i + 2])
            else:
                a.bg = int(parts[i + 2])
            i += 3
            continue
        elif v in (38, 48) and i + 4 < len(parts) and parts[i + 1] == "2":
            rgb = (int(parts[i + 2]), int(parts[i + 3]), int(parts[i + 4]))
            if v == 38:
                a.fg = rgb
            else:
                a.bg = rgb
            i += 5
            continue
        i += 1


def parse_ansi(text: str) -> list[list[Seg]]:
    lines: list[list[Seg]] = []
    curline: list[Seg] = []
    plain = ""
    attr = Attr()

    def emit() -> None:
        nonlocal plain
        if plain:
            curline.append(Seg(plain, replace(attr)))
            plain = ""

    for tok in ANSI_TOKEN_RE.findall(text):
        if tok.startswith("\x1b["):
            if tok.endswith("m"):
                emit()
                apply_sgr(attr, tok[2:-1])
            continue
        if tok.startswith(("\x1bP", "\x1b]")):
            continue
        if tok == "\r":
            continue
        if tok == "\n":
            emit()
            lines.append(curline)
            curline = []
            continue
        plain += tok
    emit()
    if curline:
        lines.append(curline)
    while lines and not lines[-1]:
        lines.pop()
    return lines


# ---------------------------------------------------------------- themes


def parse_color(s: str) -> Color | None:
    s = s.strip()
    if s.startswith("#") and len(s) == 7:
        try:
            return (int(s[1:3], 16), int(s[3:5], 16), int(s[5:7], 16))
        except ValueError:
            return None
    if s.startswith("#") and len(s) == 4:
        try:
            return (int(s[1] * 2, 16), int(s[2] * 2, 16), int(s[3] * 2, 16))
        except ValueError:
            return None
    if s.startswith("rgb:"):
        parts = s[4:].split("/")
        if len(parts) == 3 and all(parts):
            try:
                return (
                    int(parts[0][:2], 16),
                    int(parts[1][:2], 16),
                    int(parts[2][:2], 16),
                )
            except ValueError:
                return None
    return None


def kitty_config_dir() -> Path:
    k = os.environ.get("KITTY_CONFIG_DIRECTORY")
    return Path(k).expanduser() if k else Path.home() / ".config" / "kitty"


def load_kitty_conf(theme: Theme, path: Path) -> None:
    for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        key, _, val = line.partition(" ")
        if not val:
            continue
        key, val = key.strip(), val.strip()
        if key == "background":
            c = parse_color(val)
            if c is not None:
                theme.bg = c
        elif key == "foreground":
            c = parse_color(val)
            if c is not None:
                theme.fg = c
        elif key.startswith("color"):
            try:
                idx = int(key[5:])
            except ValueError:
                continue
            c = parse_color(val)
            if c is None:
                continue
            if idx < 16:
                theme.colors[idx] = c
            else:
                theme.extended[idx] = c


def find_theme_path(name: str) -> Path:
    p = Path(name)
    if p.is_file():
        return p
    dirs = [
        kitty_config_dir() / "themes",
        kitty_config_dir(),
        Path.home() / ".local/share/kitty/themes",
        Path("/usr/share/kitty/themes"),
    ]
    for d in dirs:
        for cand in (p, Path(name + ".conf")):
            f = d / cand
            if f.is_file():
                return f
    sys.exit(f"kitty theme not found: {name}")


def probe_terminal_colors() -> tuple[RGB | None, RGB | None]:
    if not (sys.stdin.isatty() and sys.stdout.isatty()):
        return None, None
    fd = sys.stdin.fileno()
    old = termios.tcgetattr(fd)
    tty.setcbreak(fd)
    buf = b""
    try:
        sys.stdout.write("\x1b]10;?\x07\x1b]11;?\x07")
        sys.stdout.flush()
        end = time.monotonic() + OSC_PROBE_TIMEOUT
        while time.monotonic() < end:
            r, _, _ = select.select([fd], [], [], max(0.0, end - time.monotonic()))
            if not r:
                break
            chunk = os.read(fd, 1024)
            if not chunk:
                break
            buf += chunk
    finally:
        termios.tcsetattr(fd, termios.TCSADRAIN, old)
    fg = bg = None
    m = re.search(rb"\x1b]10;rgb:([0-9a-f]+)/([0-9a-f]+)/([0-9a-f]+)", buf)
    if m:
        fg = (int(m.group(1)[:2], 16), int(m.group(2)[:2], 16), int(m.group(3)[:2], 16))
    m = re.search(rb"\x1b]11;rgb:([0-9a-f]+)/([0-9a-f]+)/([0-9a-f]+)", buf)
    if m:
        bg = (int(m.group(1)[:2], 16), int(m.group(2)[:2], 16), int(m.group(3)[:2], 16))
    return fg, bg


def build_theme(name: str | None, bg: str | None) -> Theme:
    theme = Theme()
    if name:
        load_kitty_conf(theme, find_theme_path(name))
    else:
        conf = kitty_config_dir() / "kitty.conf"
        if conf.is_file():
            load_kitty_conf(theme, conf)
        fg, bgc = probe_terminal_colors()
        if fg is not None:
            theme.fg = fg
        if bgc is not None:
            theme.bg = bgc
    if bg:
        c = parse_color(bg)
        if c is None:
            sys.exit(f"invalid background color: {bg}")
        theme.bg = c
    return theme


# ---------------------------------------------------------------- terminal

# Multiplexers fork new sessions exactly like a terminal would, so the
# session boundary alone cannot tell them apart from a real terminal.
PTY_MUX = {"tmux", "screen", "zellij"}

NAME_RE = re.compile(r"[A-Za-z0-9._+-]{1,40}")


def _proc_info(pid: int) -> tuple[str, int, int] | None:
    try:
        with open(f"/proc/{pid}/stat") as f:
            parts = f.read().split()
        comm = parts[1].strip("()")
        ppid = int(parts[3])
        with open(f"/proc/{pid}/sessionid") as f:
            session = int(f.read())
        return comm, ppid, session
    except (OSError, ValueError, IndexError):
        return None


def detect_terminal() -> str:
    if os.environ.get("KITTY_WINDOW_ID"):
        return "kitty"
    prog = os.environ.get("TERM_PROGRAM")
    if prog:
        return prog
    sid = os.getsid(0)
    pid = os.getppid()
    for _ in range(SCAN_DEPTH):
        if pid <= 1:
            break
        info = _proc_info(pid)
        if info is None:
            break
        comm, ppid, psid = info
        if psid != sid:
            if comm in PTY_MUX:
                pid = ppid
                continue
            return comm
        pid = ppid
    return "terminal"


TERM_SHELL_SCRIPT = '#!/bin/sh\n"${SHELL:-/bin/bash}" "$1" "$2"\n:\n'


def term_shell(name: str) -> str | None:
    if not NAME_RE.fullmatch(name):
        return None
    d = CACHE_DIR / "w"
    try:
        d.mkdir(parents=True, exist_ok=True)
        link = d / name
        link.unlink(missing_ok=True)
        link.write_text(TERM_SHELL_SCRIPT)
        link.chmod(0o755)
        return str(link)
    except OSError:
        return None


def term_alias(name: str, site: str | None) -> str | None:
    if not name or not NAME_RE.fullmatch(name):
        return None
    try:
        exe = os.readlink("/proc/self/exe")
        root = Path(exe).resolve().parent.parent
        tree = CACHE_DIR / "uvpy"
        if not tree.is_dir():
            shutil.copytree(root, tree, symlinks=True, ignore_dangling_symlinks=True)
        target = tree / "bin" / name
        if not target.is_file():
            if target.exists():
                target.unlink()
            shutil.copy2(exe, target)
            target.chmod(0o755)
        cfg = tree / "pyvenv.cfg"
        if not cfg.is_file():
            cfg.write_text(f"home = {Path(exe).resolve().parent}\n")
        if site:
            sp = (
                tree
                / "lib"
                / f"python{sys.version_info[0]}.{sys.version_info[1]}"
                / "site-packages"
            )
            sp.parent.mkdir(parents=True, exist_ok=True)
            if sp.is_symlink() and os.readlink(sp) == site:
                pass
            else:
                if sp.exists():
                    shutil.rmtree(sp)
                sp.symlink_to(site, target_is_directory=True)
        return str(target)
    except OSError:
        return None


def terminal_version(name: str) -> tuple[int, int, int] | None:
    bin_ = shutil.which(name)
    if not bin_:
        return None
    try:
        run = subprocess.run(
            [bin_, "--version"], check=False, capture_output=True, text=True, timeout=5
        )
    except (OSError, subprocess.TimeoutExpired):
        return None
    m = re.search(r"(\d+)\.(\d+)\.(\d+)", run.stdout or "")
    if not m:
        return None
    return (
        min(99, int(m.group(1))),
        min(99, int(m.group(2))),
        min(99, int(m.group(3))),
    )


# ---------------------------------------------------------------- pty capture


def kitty_terminal_font() -> tuple[str, str]:
    family = size = None
    conf = kitty_config_dir() / "kitty.conf"
    if conf.is_file():
        for line in conf.read_text(encoding="utf-8", errors="replace").splitlines():
            line = line.strip()
            if family is None and line.startswith("font_family "):
                family = line.split(None, 1)[1].strip()
            elif size is None and line.startswith("font_size "):
                size = line.split(None, 1)[1].strip()
    return family or "monospace", size or "11"


@final
class QueryResponder:
    """Answers terminal protocol queries fastfetch emits into the pty."""

    def __init__(
        self,
        fd: int,
        font: tuple[str, str],
        da_version: tuple[int, int, int] | None,
    ) -> None:
        self._fd = fd
        self._font = font
        self._da_version = da_version
        self._pos = 0

    def feed(self, buf: bytes) -> None:
        family, size = self._font
        searched = self._pos
        while True:
            di = buf.find(b"\x1bP+q", searched)
            ci = buf.find(b"\x1b[>c", searched)
            if di < 0 and ci < 0:
                self._pos = max(self._pos, len(buf) - 8)
                return
            if ci < 0 or (di >= 0 and di < ci):
                end = buf.find(b"\x1b\\", di + 4)
                if end < 0:
                    self._pos = max(self._pos, di)
                    return
                resp = b""
                for qid in buf[di + 4 : end].decode("ascii", "replace").split(";"):
                    try:
                        dec = bytes.fromhex(qid).decode("utf-8", "replace")
                    except ValueError:
                        continue
                    if dec == "kitty-query-font_family":
                        resp += (
                            b"\x1bP1+r"
                            + (dec + "=" + family).encode().hex().encode()
                            + b"\x1b\\"
                        )
                    elif dec == "kitty-query-font_size":
                        resp += (
                            b"\x1bP1+r"
                            + (dec + "=" + size).encode().hex().encode()
                            + b"\x1b\\"
                        )
                if resp:
                    try:
                        os.write(self._fd, resp)
                    except OSError:
                        pass
                searched = end + 2
            else:
                if self._da_version:
                    ver = "".join(f"{v:02d}" for v in self._da_version)
                    try:
                        os.write(self._fd, b"\x1b[>1;" + ver.encode() + b";0c")
                    except OSError:
                        pass
                searched = ci + 4


def read_pty(
    fd: int, responder: QueryResponder | None, quiet: float, deadline: float
) -> bytes:
    buf = b""
    tail = time.monotonic()
    while time.monotonic() < deadline:
        r, _, _ = select.select([fd], [], [], 0.05)
        if not r:
            if buf and time.monotonic() - tail > quiet:
                break
            continue
        try:
            chunk = os.read(fd, 65536)
        except OSError:
            break
        if not chunk:
            break
        buf += chunk
        tail = time.monotonic()
        if responder:
            responder.feed(buf)
    return buf


def reap_child(pid: int, fd: int) -> None:
    try:
        os.kill(pid, signal.SIGKILL)
        os.waitpid(pid, 0)
    except OSError:
        pass
    try:
        os.close(fd)
    except OSError:
        pass


def capture_prompt() -> str | None:
    shell = os.environ.get("SHELL") or "bash"
    if not shell or not Path(shell).is_file():
        return None
    try:
        pid, fd = pty.fork()
    except OSError:
        return None
    if pid == 0:
        try:
            os.environ["TERM"] = os.environ.get("TERM") or "xterm-256color"
            os.execv(shell, [shell, "-i"])
        except OSError:
            os._exit(127)
    buf = read_pty(
        fd, None, quiet=PROMPT_QUIET, deadline=time.monotonic() + PROMPT_DEADLINE
    )
    reap_child(pid, fd)
    if not buf:
        return None
    text = buf.replace(b"\r", b"").decode("utf-8", "replace")
    text = re.sub(r"\x1b\][^\x1b\x07]*(\x07|\x1b\\)", "", text)
    lines = text.split("\n")
    clean = [re.sub(r"\x1b\[[0-9;?]*[A-Za-z]", "", ln) for ln in lines]
    for i in range(len(clean) - 1, -1, -1):
        if clean[i].strip():
            return lines[i]
    return None


def pty_capture(terminal: str, argv: list[str]) -> str | None:
    if not os.environ.get(ALIAS_ENV):
        try:
            own = os.path.basename(os.readlink("/proc/self/exe"))
        except OSError:
            own = terminal
        if own != terminal:
            sps = [p for p in sys.path if p.endswith("site-packages")]
            alias = term_alias(terminal, sps[0] if sps else None)
            if alias:
                os.environ[ALIAS_ENV] = "1"
                os.environ["PYTHONHOME"] = str(Path(alias).resolve().parent.parent)
                try:
                    os.execv(alias, [alias] + sys.argv)
                except OSError:
                    pass
    if terminal != "kitty":
        os.environ.pop("KITTY_WINDOW_ID", None)
    try:
        pid, fd = pty.fork()
    except OSError:
        return None
    if pid == 0:
        try:
            os.environ.setdefault("KITTY_INSTALLATION_DIR", "kitty")
            if terminal == "kitty":
                os.environ.setdefault("KITTY_WINDOW_ID", "1")
            os.execv(argv[0], argv)
        except OSError:
            os._exit(127)
    try:
        attrs = termios.tcgetattr(fd)
        attrs[3] = attrs[3] & ~termios.ECHO
        termios.tcsetattr(fd, termios.TCSANOW, attrs)
    except (termios.error, OSError):
        pass
    responder = QueryResponder(
        fd,
        font=kitty_terminal_font(),
        da_version=terminal_version(terminal) if terminal == "foot" else None,
    )
    buf = read_pty(
        fd, responder, quiet=CAPTURE_QUIET, deadline=time.monotonic() + CAPTURE_DEADLINE
    )
    reap_child(pid, fd)
    return buf.replace(b"\r", b"").decode("utf-8", "replace")


def fetch_output(fetch: str, icon: str | None, raw: bool) -> str:
    cmd = [os.path.expanduser(p) for p in shlex.split(fetch)] or ["fastfetch"]
    if icon:
        cmd += ["--logo", icon]
    joined = shlex.join(cmd)
    shell = os.environ.get("SHELL") or "bash"
    inner = f"stty cols {CAPTURE_COLS}; " + joined + " ; :"
    terminal = detect_terminal()
    wrapper = term_shell(terminal)
    if wrapper:
        argv = [wrapper] + (["-ic", inner] if not raw else ["-c", inner])
        out = pty_capture(terminal, argv)
        if out:
            return out
    argv = [shell] + (["-ic", inner] if not raw else ["-c", inner])
    out = pty_capture(terminal, argv)
    if out:
        return out
    sc = shutil.which("script")
    if sc:
        run = subprocess.run(
            [sc, "-qec", inner, "/dev/null"], check=False, capture_output=True
        )
        data = run.stdout.replace(b"\r", b"")
        if data.strip():
            return data.decode("utf-8", "replace")
    run = subprocess.run(cmd, check=False, capture_output=True)
    data = run.stdout.replace(b"\r", b"")
    if not data.strip():
        sys.exit(
            "fastfetch produced no output: "
            + run.stderr.decode("utf-8", "replace").strip()
        )
    return data.decode("utf-8", "replace")


# ---------------------------------------------------------------- fonts

FONT_NAMES = ["DejaVu Sans Mono", "Liberation Mono", "Noto Sans Mono", "monospace"]
FONT_PATHS = [
    "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf",
    "/usr/share/fonts/TTF/DejaVuSansMono.ttf",
    "/usr/share/fonts/liberation/LiberationMono-Regular.ttf",
    "/usr/share/fonts/noto/NotoSansMono-Regular.ttf",
]
BOLD_FONT_PATHS = [
    "/usr/share/fonts/truetype/dejavu/DejaVuSansMono-Bold.ttf",
    "/usr/share/fonts/TTF/DejaVuSansMono-Bold.ttf",
    "/usr/share/fonts/liberation/LiberationMono-Bold.ttf",
    "/usr/share/fonts/noto/NotoSansMono-Bold.ttf",
]


@dataclass
class FontSet:
    base: ImageFont.FreeTypeFont
    bold: ImageFont.FreeTypeFont | None
    base_path: str | None
    bold_path: str | None
    fallbacks: list[str]


def scan_font_dirs() -> list[str]:
    found: list[str] = []
    dirs = [
        Path.home() / ".local/share/fonts",
        Path.home() / ".fonts",
        Path("/usr/share/fonts"),
        Path("/usr/local/share/fonts"),
    ]
    for d in dirs:
        if not d.is_dir():
            continue
        try:
            for f in sorted(d.rglob("*")):
                if f.suffix.lower() in (".ttf", ".otf"):
                    found.append(str(f))
        except OSError:
            continue
    return found


def fc_match(pattern: str) -> str | None:
    try:
        out = subprocess.run(
            ["fc-match", "-f", "%{file}", pattern],
            check=False,
            capture_output=True,
            text=True,
            timeout=10,
        )
        got = out.stdout.strip()
        return got or None
    except (OSError, subprocess.TimeoutExpired):
        return None


def charset_ranges(path: str | None) -> list[tuple[int, int]] | None:
    if not path:
        return None
    try:
        out = subprocess.run(
            ["fc-query", "--format", "%{charset}", path],
            check=False,
            capture_output=True,
            text=True,
            timeout=10,
        )
    except (OSError, subprocess.TimeoutExpired):
        return None
    ranges = []
    for part in out.stdout.strip().split():
        if "-" in part:
            a, b = part.split("-", 1)
            ranges.append((int(a, 16), int(b, 16)))
        else:
            ranges.append((int(part, 16), int(part, 16)))
    return ranges or None


def covers(ranges: list[tuple[int, int]] | None, ch: str) -> bool:
    if not ranges:
        return True
    cp = ord(ch)
    lo, hi = 0, len(ranges) - 1
    while lo <= hi:
        mid = (lo + hi) // 2
        a, b = ranges[mid]
        if cp < a:
            hi = mid - 1
        elif cp > b:
            lo = mid + 1
        else:
            return True
    return False


def try_font(
    cands: list[str], size: float
) -> tuple[ImageFont.FreeTypeFont | None, str | None]:
    for c in cands:
        try:
            return ImageFont.truetype(c, size), c
        except OSError:
            continue
    return None, None


def font_path_of(name: str | None) -> str | None:
    if name and Path(name).suffix.lower() in (".ttf", ".otf"):
        return name
    return fc_match(name) if name else None


def resolve_fonts(size: float, path: str | None) -> FontSet:
    scanned = scan_font_dirs()
    nerds = sorted(
        f for f in scanned if "Nerd" in Path(f).name or "NerdFont" in Path(f).name
    )
    nerd_mono = [f for f in nerds if "Mono" in f]
    nerd_rest = [f for f in nerds if "Mono" not in f]
    others = [f for f in scanned if f not in nerds]
    explicit = []
    if path:
        if Path(path).is_file():
            explicit.append(path)
        else:
            f = fc_match(path)
            if f:
                explicit.append(f)
    cands = explicit + nerd_mono + nerd_rest
    for pat in ("monospace", "mono"):
        m = fc_match(pat)
        if m and m not in cands:
            cands.append(m)
    cands += FONT_PATHS + FONT_NAMES + others
    base, base_name = try_font(cands, size)
    if base is None:
        sys.exit("cannot load a font, pass --font <path.ttf>")
    base_path = font_path_of(base_name)
    bold_cands = []
    if base_path:
        p = Path(base_path)
        stem = p.stem.replace("-Regular", "").replace("Regular", "")
        for cand in (f"{stem}-Bold", f"{stem}Bold", p.stem + "-Bold", p.stem + "Bold"):
            bold_cands.append(str(p.with_name(cand + p.suffix)))
    for pat in ("monospace:bold", "mono:bold"):
        m = fc_match(pat)
        if m and m not in bold_cands:
            bold_cands.append(m)
    bold_cands += BOLD_FONT_PATHS
    bold, bold_name = try_font(bold_cands, size)
    bold_path = font_path_of(bold_name) if bold is not None else None
    fallbacks: list[str] = []
    for fam in (
        "Symbols Nerd Font",
        "JetBrainsMono Nerd Font",
        "FiraCode Nerd Font",
        "Noto Sans Symbols 2",
        "Noto Sans Symbols",
        "Symbola",
        "Unifont",
    ):
        f = fc_match(fam)
        if f and f not in fallbacks:
            fallbacks.append(f)
    fallbacks += nerds + others + FONT_PATHS
    dedup: list[str] = []
    for f in fallbacks:
        if f not in dedup:
            dedup.append(f)
    if base_path:
        dedup = [f for f in dedup if f != base_path]
    return FontSet(base, bold, base_path, bold_path, dedup)


@final
class FontPicker:
    """Picks a glyph-matching font per character, with lazy caches."""

    def __init__(self, fonts: FontSet, size: float) -> None:
        self.base = fonts.base
        self.bold = fonts.bold
        self.size = size
        self.fallbacks = fonts.fallbacks
        self.primary_cs = charset_ranges(fonts.base_path)
        self.bold_cs = charset_ranges(fonts.bold_path) or self.primary_cs
        self._fonts: dict[str, ImageFont.FreeTypeFont] = {}
        self._cs: dict[str, list[tuple[int, int]] | None] = {}
        if fonts.base_path:
            self._cs[fonts.base_path] = self.primary_cs

    def covers(self, ch: str, is_bold: bool) -> bool:
        return covers(self.bold_cs if is_bold else self.primary_cs, ch)

    def pick(self, ch: str, is_bold: bool) -> ImageFont.FreeTypeFont:
        if self.bold is not None and is_bold and covers(self.bold_cs, ch):
            return self.bold
        if covers(self.primary_cs, ch):
            return self.base
        fb = self._fallback(ch)
        return fb or (self.bold if (self.bold is not None and is_bold) else self.base)

    def _fallback(self, ch: str) -> ImageFont.FreeTypeFont | None:
        base_path = next((p for p in self._cs if p), None)
        m = fc_match(f":charset={ord(ch):X}")
        cands = ([m] if m and m != base_path else []) + self.fallbacks
        for cand in cands:
            if cand not in self._cs:
                self._cs[cand] = charset_ranges(cand)
            if not covers(self._cs[cand], ch):
                continue
            f = self._fonts.get(cand)
            if f is None:
                f = try_font([cand], self.size)[0]
                if f is None:
                    continue
                self._fonts[cand] = f
            return f
        return None


# ---------------------------------------------------------------- rendering


def resolve_color(c: Color | None, bold: bool, theme: Theme) -> RGB:
    if c is None:
        return theme.foreground
    if isinstance(c, tuple):
        return c
    idx = c
    if bold and idx < 8:
        idx += 8
    if idx < 16:
        return cast(RGB, theme.colors[idx] or STANDARD[idx])
    return cast(RGB, theme.extended.get(idx) or xterm256(idx))


def render(
    lines: list[list[Seg]],
    fonts: FontSet,
    theme: Theme,
    pad: int,
    size: float,
) -> Image.Image:
    base = fonts.base
    metrics = base.getmetrics()
    line_h = metrics[0] + metrics[1] + 2
    widths = []
    for line in lines:
        w = 0.0
        for seg in line:
            f = fonts.bold if (seg.attr.bold and fonts.bold is not None) else base
            w += f.getlength(seg.text)
        widths.append(w)
    picker = FontPicker(fonts, size)
    img = Image.new(
        "RGB",
        (int(max(widths)) + 1 + 2 * pad, line_h * len(lines) + 2 * pad),
        theme.background,
    )
    d = ImageDraw.Draw(img)
    for li, line in enumerate(lines):
        top = pad + li * line_h
        baseline = pad + li * line_h + metrics[0] + 1
        x = float(pad)
        for seg in line:
            a = seg.attr
            is_bold = a.bold and fonts.bold is not None
            f = fonts.bold if is_bold else base
            if f is None:
                f = base
            fg = resolve_color(a.fg, a.bold, theme)
            bgc = None if a.bg is None else resolve_color(a.bg, False, theme)
            if a.inverse:
                fg, bgc = bgc, fg
                if fg is None:
                    fg = theme.background
            need_fb = any(not picker.covers(ch, is_bold) for ch in seg.text)
            if not need_fb:
                w = f.getlength(seg.text)
                if w <= 0:
                    continue
                if bgc is not None:
                    d.rectangle([x, top, x + w, top + line_h], fill=bgc)
                if a.bold and fonts.bold is None:
                    d.text(
                        (x, baseline),
                        seg.text,
                        font=base,
                        fill=fg,
                        anchor="ls",
                        stroke_width=1,
                    )
                d.text((x, baseline), seg.text, font=f, fill=fg, anchor="ls")
                x += w
            else:
                wl = []
                for ch in seg.text:
                    ff = picker.pick(ch, is_bold)
                    wl.append((ff, ff.getlength(ch)))
                total = sum(w for _, w in wl)
                if bgc is not None:
                    d.rectangle([x, top, x + total, top + line_h], fill=bgc)
                xc = x
                for ch, (ff, w) in zip(seg.text, wl):
                    if a.bold and fonts.bold is None and ff is base:
                        d.text(
                            (xc, baseline),
                            ch,
                            font=base,
                            fill=fg,
                            anchor="ls",
                            stroke_width=1,
                        )
                    d.text((xc, baseline), ch, font=ff, fill=fg, anchor="ls")
                    xc += w
                x += total
    return img


def save_image(img: Image.Image, path: str, fmt: str) -> None:
    if fmt == "gif":
        img.convert("P", palette=Image.ADAPTIVE, colors=256).save(path, format="GIF")  # pyright: ignore[reportAttributeAccessIssue]
    elif fmt == "jpg":
        img.save(path, format="JPEG", quality=92)
    elif fmt == "webp":
        img.save(path, format="WEBP", lossless=True)
    else:
        img.save(path, format="PNG")


# ---------------------------------------------------------------- prompt


@dataclass
class PromptParts:
    user: str
    host: str
    dir: str
    marker: str
    style: str


def prompt_parts(shell: str | None) -> PromptParts:
    try:
        user = getpass.getuser()
    except OSError:
        user = "user"
    try:
        host = socket.gethostname().split(".")[0]
    except OSError:
        host = "localhost"
    home = os.path.expanduser("~")
    cwd = os.getcwd()
    if cwd == home:
        dir_ = "~"
    elif cwd.startswith(home + os.sep):
        dir_ = "~" + cwd[len(home) :]
    else:
        dir_ = cwd
    style = (shell or os.environ.get("SHELL", "") or "bash").strip().lower()
    style = style.rsplit("/", 1)[-1]
    marker = "#" if user == "root" else "$"
    if style.startswith("fish"):
        return PromptParts(user, host, dir_, ">", "plain")
    if style.startswith("zsh"):
        return PromptParts(user, host, dir_, "%", "plain")
    if style.startswith("nu"):
        return PromptParts("", "", dir_, ">", "nu")
    if style.startswith("pwsh") or "powershell" in style:
        return PromptParts("", "", dir_, ">", "pwsh")
    return PromptParts(user, host, dir_, marker, "bracket")


def make_prompt(shell: str | None) -> list[Seg]:
    p = prompt_parts(shell)

    def bright(i: int) -> Attr:
        return Attr(fg=i + 8, bold=True)

    if p.style == "nu":
        return [Seg(p.dir, bright(4)), Seg("> ", Attr())]
    if p.style == "pwsh":
        return [Seg("PS ", bright(4)), Seg(p.dir, Attr()), Seg("> ", Attr())]
    if p.style == "plain":
        return [
            Seg(f"{p.user}@{p.host}", bright(2)),
            Seg(f" {p.dir}", bright(4)),
            Seg(p.marker + " ", Attr()),
        ]
    return [
        Seg("[", Attr()),
        Seg(f"{p.user}@{p.host}", bright(2)),
        Seg(f" {p.dir}]", bright(4)),
        Seg(p.marker + " ", Attr()),
    ]


def command_line(fetch: str, icon: str | None) -> list[Seg]:
    cmd = [fetch]
    if icon:
        cmd += ["--logo", icon]
    return [Seg(shlex.join(cmd), Attr())]


def prompt_lines(prompt_text: str | None) -> list[list[Seg]] | None:
    if not prompt_text:
        return None
    lines = parse_ansi(prompt_text)
    while lines and not lines[-1]:
        lines.pop()
    return lines or None


# ---------------------------------------------------------------- cli


class Args(argparse.Namespace):
    output: str | None
    format: str | None
    icon: str | None
    theme: str | None
    prompt: str | None
    shell: str | None
    raw: bool
    bg: str | None
    scale: float
    font_size: int
    font: str | None
    padding: int
    fastfetch: str


def parse_args(argv: list[str] | None) -> Args:
    p = argparse.ArgumentParser(
        prog="fetch-snap",
        description="Render fastfetch output to an image, keeping all colors.",
    )
    p.add_argument(
        "-o",
        "--output",
        help="output file (default: fetch.<format>); extension picks the format",
    )
    p.add_argument(
        "-f",
        "--format",
        choices=["png", "jpg", "webp", "gif"],
        help="image format (default: png)",
    )
    p.add_argument(
        "--icon",
        help="fetch icon (logo) to use, passed to fastfetch --logo, e.g. ascii, nerd_font",
    )
    p.add_argument(
        "--theme",
        help="kitty color theme name (searched in kitty theme dirs) or a .conf file path",
    )
    p.add_argument(
        "--prompt",
        choices=["top", "bottom", "both"],
        help="add a real shell prompt line above and/or below the output, captured from $SHELL",
    )
    p.add_argument(
        "--shell",
        help="fallback prompt style if the real prompt cannot be captured: bash, zsh, fish, nu, pwsh (default: $SHELL)",
    )
    p.add_argument(
        "--raw",
        action="store_true",
        help="run the fastfetch command directly, do not expand shell aliases",
    )
    p.add_argument(
        "--bg",
        help="background color, e.g. #1e1e2e (default: dark gray, or the theme background)",
    )
    p.add_argument(
        "--scale",
        type=float,
        default=2,
        help="render scale, multiplies font size and padding (default: 2)",
    )
    p.add_argument(
        "--font-size",
        type=int,
        default=16,
        help="font size in px at scale 1 (default: 16)",
    )
    p.add_argument("--font", help="path to a ttf/otf font")
    p.add_argument(
        "--padding",
        type=int,
        default=14,
        help="padding around the text in px (default: 14)",
    )
    p.add_argument(
        "--fastfetch",
        default="fastfetch",
        help="command to run (default: fastfetch), run through an interactive shell unless --raw",
    )
    return cast(Args, p.parse_args(argv))


def resolve_output(output: str | None, fmt: str | None) -> tuple[str, str]:
    if not fmt:
        ext = Path(output).suffix[1:].lower() if output else ""
        fmt = ext if ext in ("png", "jpg", "jpeg", "webp", "gif") else "png"
    if fmt == "jpeg":
        fmt = "jpg"
    return output or f"fetch.{fmt}", fmt


def main(argv: list[str] | None = None) -> None:
    args = parse_args(argv)
    theme = build_theme(args.theme, args.bg)
    text = fetch_output(args.fastfetch, args.icon, args.raw)
    lines = parse_ansi(text)
    prompt = prompt_lines(capture_prompt() if args.prompt else None)
    if args.prompt in ("top", "both"):
        top = list(prompt or [make_prompt(args.shell)])
        top[-1] = top[-1] + command_line(args.fastfetch, args.icon)
        lines[:0] = top
    if args.prompt in ("bottom", "both"):
        lines.extend(prompt or [make_prompt(args.shell)])
    size = args.font_size * args.scale
    pad = round(args.padding * args.scale)
    fonts = resolve_fonts(size, args.font)
    if fonts.base_path:
        print(f"font: {fonts.base_path}", file=sys.stderr)
    if fonts.bold_path and fonts.bold_path != fonts.base_path:
        print(f"bold: {fonts.bold_path}", file=sys.stderr)
    img = render(lines, fonts, theme, pad, size)
    out, fmt = resolve_output(args.output, args.format)
    save_image(img, out, fmt)
    print(f"wrote {out}")


if __name__ == "__main__":
    main()