#!/usr/bin/env python3
"""Render two terminal pane captures (with ANSI escapes) into a single PNG.

Usage:
    python3 scripts/render_terminal.py left.ansi right.ansi -o assets/screenshot.png
    python3 scripts/render_terminal.py left.ansi right.ansi --labels "alice" "bob" -o out.png
"""

import argparse
import re
import sys
from pathlib import Path

from PIL import Image, ImageDraw, ImageFont

# ── Theme (dark terminal) ────────────────────────────────────────────────────
BG        = (30, 30, 46)      # base (catppuccin mocha-ish)
FG        = (205, 214, 244)   # default text
DIM       = (108, 112, 134)   # dim/grey
GREEN     = (166, 227, 161)
CYAN      = (137, 220, 235)
YELLOW    = (249, 226, 175)
RED       = (243, 139, 168)
BLUE      = (137, 180, 250)
MAGENTA   = (203, 166, 247)
BOLD_WHITE = (255, 255, 255)
BORDER    = (69, 71, 90)
TITLE_BG  = (49, 50, 68)

ANSI_COLORS = {
    30: (30, 30, 46),     31: RED,          32: GREEN,    33: YELLOW,
    34: BLUE,             35: MAGENTA,      36: CYAN,     37: FG,
    90: DIM,              91: RED,          92: GREEN,    93: YELLOW,
    94: BLUE,             95: MAGENTA,      96: CYAN,     97: BOLD_WHITE,
}

# ── ANSI parsing ─────────────────────────────────────────────────────────────
ESC_RE = re.compile(r'\x1b\[([0-9;]*)m')

def parse_ansi_line(line):
    """Yield (text, fg_color, bold) spans from an ANSI-escaped line."""
    fg = FG
    bold = False
    dim = False
    pos = 0
    for m in ESC_RE.finditer(line):
        if m.start() > pos:
            color = fg
            if dim and color == FG:
                color = DIM
            yield (line[pos:m.start()], color, bold)
        codes = m.group(1).split(';') if m.group(1) else ['0']
        for code_s in codes:
            code = int(code_s) if code_s else 0
            if code == 0:
                fg, bold, dim = FG, False, False
            elif code == 1:
                bold = True
            elif code == 2:
                dim = True
            elif code in ANSI_COLORS:
                fg = ANSI_COLORS[code]
        pos = m.end()
    tail = line[pos:]
    if tail:
        color = fg
        if dim and color == FG:
            color = DIM
        yield (tail, color, bold)


def load_font(size):
    """Try to load a monospace font."""
    candidates = [
        "/usr/share/fonts/google-noto/NotoSansMono-Regular.ttf",
        "/usr/share/fonts/truetype/noto/NotoSansMono-Regular.ttf",
        "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf",
        "/usr/share/fonts/dejavu-sans-mono-fonts/DejaVuSansMono.ttf",
        "/usr/share/fonts/TTF/DejaVuSansMono.ttf",
        "/usr/share/fonts/liberation-mono/LiberationMono-Regular.ttf",
        "/usr/share/fonts/google-droid-sans-mono-fonts/DroidSansMono.ttf",
    ]
    for path in candidates:
        if Path(path).exists():
            return ImageFont.truetype(path, size)
    # fallback
    return ImageFont.load_default()


def load_bold_font(size):
    candidates = [
        "/usr/share/fonts/google-noto/NotoSansMono-Bold.ttf",
        "/usr/share/fonts/truetype/noto/NotoSansMono-Bold.ttf",
        "/usr/share/fonts/truetype/dejavu/DejaVuSansMono-Bold.ttf",
        "/usr/share/fonts/dejavu-sans-mono-fonts/DejaVuSansMono-Bold.ttf",
        "/usr/share/fonts/TTF/DejaVuSansMono-Bold.ttf",
        "/usr/share/fonts/liberation-mono/LiberationMono-Bold.ttf",
    ]
    for path in candidates:
        if Path(path).exists():
            return ImageFont.truetype(path, size)
    return load_font(size)


def strip_ansi(s):
    return ESC_RE.sub('', s)


def render_pane(lines, width_chars, font, bold_font, font_size, line_height):
    """Render terminal lines to an Image."""
    char_w = font.getbbox("M")[2]
    img_w = char_w * width_chars + 24  # 12px padding each side
    img_h = line_height * len(lines) + 16  # 8px padding top+bottom

    img = Image.new("RGB", (img_w, img_h), BG)
    draw = ImageDraw.Draw(img)

    y = 8
    for line in lines:
        x = 12
        for text, color, bold in parse_ansi_line(line):
            f = bold_font if bold else font
            draw.text((x, y), text, fill=color, font=f)
            x += f.getlength(text)
        y += line_height

    return img


def main():
    ap = argparse.ArgumentParser(description="Render terminal panes to PNG")
    ap.add_argument("left", help="Left pane ANSI capture file")
    ap.add_argument("right", help="Right pane ANSI capture file")
    ap.add_argument("-o", "--output", default="assets/screenshot.png")
    ap.add_argument("--labels", nargs=2, default=["alice", "bob"],
                    help="Labels for the two panes")
    ap.add_argument("--font-size", type=int, default=14)
    ap.add_argument("--width", type=int, default=58,
                    help="Width of each pane in characters")
    args = ap.parse_args()

    font_size = args.font_size
    line_height = int(font_size * 1.5)
    font = load_font(font_size)
    bold_font = load_bold_font(font_size)
    char_w = font.getbbox("M")[2]
    pane_w = char_w * args.width + 24

    left_lines = Path(args.left).read_text().splitlines()
    right_lines = Path(args.right).read_text().splitlines()

    # Render each pane
    left_img = render_pane(left_lines, args.width, font, bold_font,
                           font_size, line_height)
    right_img = render_pane(right_lines, args.width, font, bold_font,
                            font_size, line_height)

    # Composite: title bar + two panes side by side
    title_h = 32
    gap = 2
    max_h = max(left_img.height, right_img.height)
    total_w = left_img.width + gap + right_img.width
    total_h = title_h + max_h

    # Window chrome
    canvas = Image.new("RGB", (total_w, total_h), BORDER)
    draw = ImageDraw.Draw(canvas)

    # Title bar
    draw.rectangle([(0, 0), (total_w, title_h - 1)], fill=TITLE_BG)

    # Traffic lights
    for i, color in enumerate([(255, 95, 86), (255, 189, 46), (39, 201, 63)]):
        cx = 16 + i * 22
        cy = title_h // 2
        draw.ellipse([(cx - 6, cy - 6), (cx + 6, cy + 6)], fill=color)

    # Pane labels
    label_font = load_font(font_size - 1)
    left_label = args.labels[0]
    right_label = args.labels[1]
    left_label_w = label_font.getlength(left_label)
    right_label_w = label_font.getlength(right_label)
    draw.text((left_img.width // 2 - left_label_w // 2, 7),
              left_label, fill=DIM, font=label_font)
    draw.text((left_img.width + gap + right_img.width // 2 - right_label_w // 2, 7),
              right_label, fill=DIM, font=label_font)

    # Paste panes
    canvas.paste(left_img, (0, title_h))
    canvas.paste(right_img, (left_img.width + gap, title_h))

    # Round corners (simple mask)
    radius = 10
    mask = Image.new("L", canvas.size, 255)
    mask_draw = ImageDraw.Draw(mask)
    mask_draw.rectangle([(0, 0), (radius, radius)], fill=0)
    mask_draw.pieslice([(0, 0), (radius * 2, radius * 2)], 180, 270, fill=255)
    mask_draw.rectangle([(total_w - radius, 0), (total_w, radius)], fill=0)
    mask_draw.pieslice([(total_w - radius * 2, 0), (total_w, radius * 2)], 270, 360, fill=255)

    # Apply rounded corners with transparent background
    final = Image.new("RGBA", canvas.size, (0, 0, 0, 0))
    final.paste(canvas, mask=mask)

    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    final.save(args.output)
    print(f"Saved {args.output} ({final.width}x{final.height})")


if __name__ == "__main__":
    main()
