"""Generate a complete Matrix character grid as square SVG and PNG avatars.

Run with: uv run python cmatrix_avatar.py --seed 42 --previews
"""

from __future__ import annotations

import argparse
from contextlib import ExitStack
from dataclasses import dataclass
import os
from pathlib import Path
import random
import secrets
from xml.sax.saxutils import escape

from fontTools.pens.boundsPen import BoundsPen
from fontTools.pens.svgPathPen import SVGPathPen
from fontTools.pens.transformPen import TransformPen
from fontTools.ttLib import TTFont
from PIL import Image
import resvg_py
from avatar_color import gradient_level, sdr_color, validate_peak


BACKGROUND = "#0b0b0b"
BLUE = "#3b73fa"
YELLOW = "#c7a52c"
GREEN = "#00ff00"
GWB = (
    " ██████╗ ██╗    ██╗██████╗ ",
    "██╔════╝ ██║    ██║██╔══██╗",
    "██║  ███╗██║ █╗ ██║██████╔╝",
    "██║   ██║██║███╗██║██╔══██╗",
    "╚██████╔╝╚███╔███╔╝██████╔╝",
    " ╚═════╝  ╚══╝╚══╝ ╚═════╝ ",
)
QUOTES = (
    "Never Lose a Holy Curiosity.",
    "What I cannot create, I do not understand.",
    "A person who never made a mistake never tried anything new.",
    "Not even wrong.",
    "For a successful technology, reality must take precedence over public relations, for nature cannot be fooled.",
    "Know how to solve every problem that has been solved.",
    "I would rather have questions that can't be answered than answers that can't be questioned.",
    "We are trying to prove ourselves wrong as quickly as possible, because only in that way can we find progress.",
    "Education is what remains after one has forgotten everything one has learned in school.",
    "The important thing is not to stop questioning. Curiosity has its own reason for existence. One cannot help but be in awe when he contemplates the mysteries of eternity, of life, of the marvelous structure of reality. It is enough if one tries merely to comprehend a little of this mystery each day.",
    "I owe a lot to my engineering training because it taught me to tolerateapproximations. Previously to that I thought one should just concentrate on exactequations all the time. Then I got the idea that in the actual world all our equations areonly approximate. We must just tend to greater and greater accuracy. In spite of theequations being approximate, they can be beautiful.",
)


@dataclass(frozen=True)
class Cell:
    char: str = " "
    color: str = BACKGROUND
    level: float = 1.0


class CMatrixAvatar:
    """Original rain algorithm, with colors separate from characters."""

    def __init__(self, seed: int | None = None, rows: int = 43, columns: int = 77):
        if not len(GWB) + 2 <= rows <= 200:
            raise ValueError(f"--rows must be between {len(GWB) + 2} and 200")
        if not max(map(len, GWB)) + 2 <= columns <= 200:
            raise ValueError(f"--columns must be between {max(map(len, GWB)) + 2} and 200")
        self.rows, self.columns = rows, columns
        self.seed = secrets.randbits(32) if seed is None else seed
        rng = random.Random(self.seed)
        text = "".join(QUOTES).replace(" ", "")
        text = "".join(c.lower() if rng.randint(0, 1) == 0 else c.upper() for c in text)
        self.cells = [[Cell() for _ in range(self.columns)] for _ in range(self.rows)]
        for column in range(0, self.columns, 2):
            length = rng.randint(0, self.rows)
            start = rng.randint(0, self.rows)
            for offset in range(length):
                position = offset / (length - 1) if length > 1 else 1.0
                self.cells[(start + offset) % self.rows][column] = Cell("+", GREEN, gradient_level(position))

        # The logo is now an independent layer. Its reserved area is determined
        # during layout; the complete rain grid remains available here.
        text_index = 0
        for row in range(self.rows):
            for column in range(self.columns):
                if self.cells[row][column].char == "+":
                    char = text[text_index % len(text)]
                    color = YELLOW if char.lower() in "gwb" else self.cells[row][column].color
                    self.cells[row][column] = Cell(char, color, self.cells[row][column].level)
                    text_index += 1


def find_font(requested: Path | None = None, *, for_logo: bool = False) -> Path:
    if requested is not None:
        if not requested.is_file():
            raise ValueError(f"Font not found: {requested}")
        return requested
    fonts = Path(os.environ.get("WINDIR", "C:/Windows")) / "Fonts"
    candidates = (
        fonts / "consola.ttf", fonts / "CascadiaMono.ttf",
        Path("/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"),
        Path("/System/Library/Fonts/Menlo.ttc"),
    )
    if not for_logo:
        candidates = (fonts / "OCRAEXT.TTF",) + candidates
    for candidate in candidates:
        if candidate.is_file():
            return candidate
    raise ValueError("No supported monospace font found; pass --font /path/to/font.ttf")


def make_svg(avatar: CMatrixAvatar, font_path: Path, size: int = 2048,
             logo_font_path: Path | None = None, logo_width: float = 0.72) -> str:
    """Outline all glyphs: no external font or text rendering dependency."""
    if not 0.4 <= logo_width <= 0.85:
        raise ValueError("--logo-width must be between 0.4 and 0.85")
    with ExitStack() as stack:
        font = stack.enter_context(TTFont(font_path, fontNumber=0))
        cmap = font.getBestCmap()
        logo_characters = set("".join(GWB).replace(" ", ""))
        if logo_font_path is None:
            logo_font_path = (font_path if all(ord(c) in cmap for c in logo_characters)
                              else find_font(for_logo=True))
        logo_font = stack.enter_context(TTFont(logo_font_path, fontNumber=0))
        if ord("M") not in cmap:
            raise ValueError("Background font must contain Latin letters.")
        cell_width = font["hmtx"][cmap[ord("M")]][0]
        cell_height = cell_width * 1.8
        sources = {}
        for is_logo, source in ((False, font), (True, logo_font)):
            glyphs, mapping = source.getGlyphSet(), source.getBestCmap()
            chars = {cell.char for row in avatar.cells for cell in row
                     if cell.char != " " and (cell.color == BLUE) == is_logo}
            if is_logo:
                chars.update(logo_characters)
            missing = sorted(c for c in chars if ord(c) not in mapping)
            if missing:
                raise ValueError(f"Font is missing characters: {''.join(missing)}")
            widths = {source["hmtx"][mapping[ord(c)]][0] for c in chars}
            if is_logo and len(widths) > 1:
                raise ValueError("Please choose a logo font with equal character advances.")
            # Display fonts such as HeavyData have varying advances. Fit their
            # widest glyph to the grid, then center narrower glyphs unchanged.
            advance = max(widths, default=cell_width)
            if ord("█") in mapping:
                block_bounds = BoundsPen(glyphs)
                glyphs[mapping[ord("█")]].draw(block_bounds)
                _, bottom, _, top = block_bounds.bounds
            else:
                # OCR A Extended has no box/block glyphs. Its Latin letters use
                # its own line metrics; the logo fallback fits the SAME cells.
                top, bottom = source["hhea"].ascent, source["hhea"].descent
            sources[is_logo] = (glyphs, mapping, cell_width / advance,
                                cell_height / (top - bottom), top, source["hmtx"], BoundsPen(glyphs))
        # Measure the full background before reserving the logo area.
        background_glyphs = []
        for row, cells in enumerate(avatar.cells):
            for column, cell in enumerate(cells):
                if cell.char == " ":
                    continue
                glyphs, mapping, sx, sy, glyph_top, metrics, bounds = sources[cell.color == BLUE]
                offset_x = (cell_width - metrics[mapping[ord(cell.char)]][0] * sx) / 2
                transform = (sx, 0, 0, -sy, column * cell_width + offset_x,
                             row * cell_height + glyph_top * sy)
                glyph = glyphs[mapping[ord(cell.char)]]
                glyph.draw(TransformPen(bounds, transform))
                cell_bounds = BoundsPen(glyphs)
                glyph.draw(TransformPen(cell_bounds, transform))
                gx0, gy0, gx1, gy1 = cell_bounds.bounds
                occupied = (min(column * cell_width, gx0), min(row * cell_height, gy0),
                            max((column + 1) * cell_width, gx1), max((row + 1) * cell_height, gy1))
                background_glyphs.append((cell, glyph, glyphs, transform, occupied))

        # Fit the COMPLETE grid and actual glyph extents, with half a cell of
        # safety space on every side. Never crop or discard background cells.
        drawn_bounds = [source[-1].bounds for source in sources.values() if source[-1].bounds]
        bx0 = min((b[0] for b in drawn_bounds), default=0)
        by0 = min((b[1] for b in drawn_bounds), default=0)
        bx1 = max((b[2] for b in drawn_bounds), default=0)
        by1 = max((b[3] for b in drawn_bounds), default=0)
        x0, y0 = min(0, bx0), min(0, by0)
        x1 = max(avatar.columns * cell_width, bx1)
        y1 = max(avatar.rows * cell_height, by1)
        side = max(x1 - x0, y1 - y0) + cell_width
        left = (side - (x1 - x0)) / 2 - x0
        top = (side - (y1 - y0)) / 2 - y0
        # Build the original six-line GWB, then scale the WHOLE logo uniformly.
        glyphs, mapping, sx, sy, glyph_top, metrics, _ = sources[True]
        logo_pen, logo_bounds = SVGPathPen(glyphs), BoundsPen(glyphs)
        for row, line in enumerate(GWB):
            for column, char in enumerate(line):
                if char == " ":
                    continue
                transform = (sx, 0, 0, -sy, column * cell_width,
                             row * cell_height + glyph_top * sy)
                glyph = glyphs[mapping[ord(char)]]
                glyph.draw(TransformPen(logo_pen, transform))
                glyph.draw(TransformPen(logo_bounds, transform))
        lx0, ly0, lx1, ly1 = logo_bounds.bounds
        logo_scale = side * logo_width / (lx1 - lx0)
        logo_height = (ly1 - ly0) * logo_scale
        logo_x, logo_y = (side - side * logo_width) / 2, (side - logo_height) / 2
        gap = cell_width * 0.25
        reserved = (logo_x - left - gap, logo_y - top - gap,
                    logo_x + side * logo_width - left + gap, logo_y + logo_height - top + gap)
        # Leave entire intersecting cells empty. No clipped glyphs or mask over
        # partial letters; all non-overlapping rain cells are retained.
        pens = {}
        rendered = 0
        for cell, glyph, glyphs, transform, (gx0, gy0, gx1, gy1) in background_glyphs:
            if gx0 < reserved[2] and gx1 > reserved[0] and gy0 < reserved[3] and gy1 > reserved[1]:
                continue
            pen = pens.setdefault((cell.color, cell.level), SVGPathPen(glyphs))
            glyph.draw(TransformPen(pen, transform))
            rendered += 1
        parts = [
            f'<svg xmlns="http://www.w3.org/2000/svg" width="{size}" height="{size}" viewBox="0 0 {side} {side}">',
            '<title>GWB Matrix avatar</title>',
            f'<desc>Complete {avatar.rows} x {avatar.columns} grid; seed {avatar.seed}; font {escape(font_path.name)}; logo font {escape(logo_font_path.name)}</desc>',
            f'<rect width="{side}" height="{side}" fill="{BACKGROUND}"/>',
            f'<g id="background" data-visible-cells="{rendered}" transform="translate({left} {top})">',
        ]
        for (color, level), pen in pens.items():
            parts.append(f'<path fill="{sdr_color(color, level)}" data-base="{color}" data-level="{level}" d="{pen.getCommands()}"/>')
        parts.extend(['</g>',
                      f'<path id="logo" fill="{sdr_color(BLUE)}" data-base="{BLUE}" data-level="1" transform="translate({logo_x} {logo_y}) scale({logo_scale}) translate({-lx0} {-ly0})" d="{logo_pen.getCommands()}"/>',
                      '</svg>'])
        return "\n".join(parts) + "\n"


def export(avatar: CMatrixAvatar, output: Path, font: Path, size: int,
           previews: bool = False, logo_font: Path | None = None,
           logo_width: float = 0.72, hdr: bool = False,
           peak_nits: float = 400) -> tuple[Path, ...]:
    if output.suffix.lower() not in ("", ".png", ".svg"):
        raise ValueError("Output must be a basename, .png, or .svg path.")
    validate_peak(peak_nits)
    png_path, svg_path = output.with_suffix(".png"), output.with_suffix(".svg")
    svg = make_svg(avatar, font, size, logo_font, logo_width)
    # Render exactly the same outlined SVG, not a second font drawing pipeline.
    png = resvg_py.svg_to_bytes(svg_string=svg)
    png_path.parent.mkdir(parents=True, exist_ok=True)
    svg_path.write_text(svg, encoding="utf-8")
    png_path.write_bytes(png)
    if previews:
        with Image.open(png_path) as image:
            for small_size in (128, 64):
                image.resize((small_size, small_size), Image.Resampling.LANCZOS).save(
                    png_path.with_name(f"{png_path.stem}-{small_size}.png"))
    if hdr:
        from avatar_hdr import export_hdr
        hdr_path = png_path.with_name(f"{png_path.stem}-hdr.png")
        export_hdr(svg, hdr_path, peak_nits)
        return svg_path, png_path, hdr_path
    return svg_path, png_path


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=Path("avatar.png"), help="Output basename or PNG/SVG path; both formats are written")
    parser.add_argument("--size", type=int, default=2048, help="Square PNG size (default: 2048)")
    parser.add_argument("--seed", type=int, help="Seed to reproduce the same rain")
    parser.add_argument("--font", type=Path, help="Monospace TTF/OTF/TTC font")
    parser.add_argument("--logo-font", type=Path, help="Optional block/box font (automatic fallback if needed)")
    parser.add_argument("--rows", type=int, default=43, help="Complete grid rows (default: 43)")
    parser.add_argument("--columns", type=int, default=77, help="Complete grid columns (default: 77)")
    parser.add_argument("--logo-width", type=float, default=0.72, help="Logo width / canvas width (default: 0.72)")
    parser.add_argument("--previews", action="store_true", help="Also save 128px and 64px previews")
    parser.add_argument("--hdr", action="store_true", help="Also export a true 16-bit BT.2020/PQ HDR PNG")
    parser.add_argument("--peak-nits", type=float, default=400, help="Shared HDR luminance target for logo and rain peaks (100..2000; default: 400)")
    args = parser.parse_args()
    if not 64 <= args.size <= 4096:
        parser.error("--size must be between 64 and 4096")
    try:
        font = find_font(args.font)
        avatar = CMatrixAvatar(args.seed, args.rows, args.columns)
        logo_font = find_font(args.logo_font, for_logo=True) if args.logo_font else None
        paths = export(avatar, args.output, font, args.size, args.previews, logo_font,
                       args.logo_width, args.hdr, args.peak_nits)
    except (OSError, ValueError) as exc:
        parser.error(str(exc))
    print(f"Seed: {avatar.seed}; font: {font.name}; PNG: {args.size} x {args.size}")
    for path in paths:
        print(f"Saved {path.resolve()}")


if __name__ == "__main__":
    main()