#!/usr/bin/python3
###############################################################################
#                                                                             #
# fireperf - A network benchmarking tool                                      #
# Copyright (C) 2026 IPFire Development Team                                  #
#                                                                             #
# This program is free software: you can redistribute it and/or modify        #
# it under the terms of the GNU General Public License as published by        #
# the Free Software Foundation, either version 3 of the License, or           #
# (at your option) any later version.                                         #
#                                                                             #
# This program is distributed in the hope that it will be useful,             #
# but WITHOUT ANY WARRANTY; without even the implied warranty of              #
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the               #
# GNU General Public License for more details.                                #
#                                                                             #
# You should have received a copy of the GNU General Public License           #
# along with this program.  If not, see <http://www.gnu.org/licenses/>.       #
#                                                                             #
###############################################################################

"""
	Turns what a terminal shows (as "tmux capture-pane -p -e" writes it) into a
	picture, which is what the screenshot in the README has been made with:

	  tmux new-session -d -x 140 -y 42 "fireperf --server"
	  ... (let it run for a while)
	  tmux capture-pane -p -e > screenshot.ansi
	  tools/ansi2png screenshot.ansi screenshot.png

	This needs Pillow and the DejaVu Sans Mono font.
"""

import argparse
import re

from PIL import Image, ImageDraw, ImageFont

# A dark theme
BACKGROUND = (0x1e, 0x1f, 0x29)
FOREGROUND = (0xd8, 0xdb, 0xe2)

# And a light one
LIGHT_BACKGROUND = (0xfa, 0xfa, 0xf7)
LIGHT_FOREGROUND = (0x2e, 0x30, 0x38)

PALETTE = [
	(0x21, 0x22, 0x2c), # black
	(0xff, 0x6e, 0x6e), # red
	(0x69, 0xff, 0x94), # green
	(0xff, 0xd8, 0x66), # yellow
	(0x6c, 0x8c, 0xff), # blue
	(0xff, 0x92, 0xdf), # magenta
	(0x7d, 0xe6, 0xff), # cyan
	(0xf8, 0xf8, 0xf2), # white
]

BRIGHT = [tuple(min(255, c + 0x30) for c in colour) for colour in PALETTE]

FONTS = "/usr/share/fonts/truetype/dejavu"

SGR = re.compile(r"\x1b\[([0-9;]*)m")

class Cell(object):
	def __init__(self, char, fg, bg, bold, dim, underline, reverse):
		self.char = char
		self.fg = fg
		self.bg = bg
		self.bold = bold
		self.dim = dim
		self.underline = underline
		self.reverse = reverse

def colour256(n):
	if n < 8:
		return PALETTE[n]
	if n < 16:
		return BRIGHT[n - 8]
	if n < 232:
		n -= 16
		return tuple(0 if v == 0 else 55 + v * 40 for v in (n // 36, n // 6 % 6, n % 6))

	v = 8 + (n - 232) * 10
	return (v, v, v)

def parse(text):
	"""
		Returns the lines of what the terminal shows as lists of cells
	"""
	lines = []

	fg, bg = None, None
	bold = dim = underline = reverse = False

	for line in text.split("\n"):
		cells = []
		pos = 0

		for m in SGR.finditer(line + "\x1b[m"):
			for char in line[pos:m.start()]:
				cells.append(Cell(char, fg, bg, bold, dim, underline, reverse))

			pos = m.end()

			codes = [int(c) if c else 0 for c in m.group(1).split(";")]

			while codes:
				c = codes.pop(0)

				if c == 0:
					fg, bg = None, None
					bold = dim = underline = reverse = False
				elif c == 1:
					bold = True
				elif c == 2:
					dim = True
				elif c == 4:
					underline = True
				elif c == 7:
					reverse = True
				elif c == 22:
					bold = dim = False
				elif c == 24:
					underline = False
				elif c == 27:
					reverse = False
				elif 30 <= c <= 37:
					fg = PALETTE[c - 30]
				elif c == 39:
					fg = None
				elif 40 <= c <= 47:
					bg = PALETTE[c - 40]
				elif c == 49:
					bg = None
				elif 90 <= c <= 97:
					fg = BRIGHT[c - 90]
				elif 100 <= c <= 107:
					bg = BRIGHT[c - 100]
				elif c in (38, 48) and codes:
					kind = codes.pop(0)

					if kind == 5 and codes:
						colour = colour256(codes.pop(0))
					elif kind == 2 and len(codes) >= 3:
						colour = tuple(codes[:3])
						codes = codes[3:]
					else:
						continue

					if c == 38:
						fg = colour
					else:
						bg = colour

		lines.append(cells)

	# Remove empty lines at the end
	while lines and not lines[-1]:
		lines.pop()

	return lines

def blend(c1, c2, f):
	return tuple(round(a * (1 - f) + b * f) for a, b in zip(c1, c2))

# Characters that we draw ourselves, so that they line up exactly: the
# fraction of the cell that they fill from the bottom or from the left
LOWER = { "▁" : 1, "▂" : 2, "▃" : 3, "▄" : 4, "▅" : 5, "▆" : 6, "▇" : 7, "█" : 8 }
LEFT  = { "▏" : 1, "▎" : 2, "▍" : 3, "▌" : 4, "▋" : 5, "▊" : 6, "▉" : 7 }

# Lines that we draw ourselves, so that they connect: which sides of the cell they
# go to (left, right, up, down), or "dotted" for a dotted horizontal line
LINES = {
	"─" : "lr", "│" : "ud", "┌" : "rd", "┐" : "ld", "└" : "ru", "┘" : "lu",
	"╭" : "rd", "╮" : "ld", "╰" : "ru", "╯" : "lu", "├" : "rud", "┤" : "lud",
	"┬" : "lrd", "┴" : "lru", "┼" : "lrud", "┈" : "dotted",
}

ROUNDED = "╭╮╰╯"

def draw_line(draw, char, x0, y0, w, h, colour, width):
	sides = LINES[char]

	# The middle of the cell
	cx = x0 + w // 2
	cy = y0 + h // 2

	if sides == "dotted":
		for x in range(x0 + 1, x0 + w, 3):
			draw.point((x, cy), fill=colour)
		return

	# Rounded corners are a quarter of a circle
	if char in ROUNDED:
		r = min(w, h) // 2

		if char == "╭":
			draw.arc((cx, cy, cx + 2 * r, cy + 2 * r), 180, 270, fill=colour, width=width)
			draw.line((cx + r, cy, x0 + w, cy), fill=colour, width=width)
			draw.line((cx, cy + r, cx, y0 + h), fill=colour, width=width)
		elif char == "╮":
			draw.arc((cx - 2 * r, cy, cx, cy + 2 * r), 270, 360, fill=colour, width=width)
			draw.line((x0, cy, cx - r, cy), fill=colour, width=width)
			draw.line((cx, cy + r, cx, y0 + h), fill=colour, width=width)
		elif char == "╰":
			draw.arc((cx, cy - 2 * r, cx + 2 * r, cy), 90, 180, fill=colour, width=width)
			draw.line((cx + r, cy, x0 + w, cy), fill=colour, width=width)
			draw.line((cx, y0, cx, cy - r), fill=colour, width=width)
		else:
			draw.arc((cx - 2 * r, cy - 2 * r, cx, cy), 0, 90, fill=colour, width=width)
			draw.line((x0, cy, cx - r, cy), fill=colour, width=width)
			draw.line((cx, y0, cx, cy - r), fill=colour, width=width)

		return

	if "l" in sides:
		draw.line((x0, cy, cx, cy), fill=colour, width=width)
	if "r" in sides:
		draw.line((cx, cy, x0 + w, cy), fill=colour, width=width)
	if "u" in sides:
		draw.line((cx, y0, cx, cy), fill=colour, width=width)
	if "d" in sides:
		draw.line((cx, cy, cx, y0 + h), fill=colour, width=width)

def render(lines, output, size, title, light=False):
	global BACKGROUND, FOREGROUND

	if light:
		BACKGROUND, FOREGROUND = LIGHT_BACKGROUND, LIGHT_FOREGROUND

	regular = ImageFont.truetype("%s/DejaVuSansMono.ttf" % FONTS, size)
	bold    = ImageFont.truetype("%s/DejaVuSansMono-Bold.ttf" % FONTS, size)

	# The size of a cell
	w = round(regular.getlength("M"))
	ascent, descent = regular.getmetrics()
	h = ascent + descent + round(size * 0.1)

	cols = max(len(l) for l in lines)
	rows = len(lines)

	# Leave space for a window around it
	margin = round(size * 1.6)
	bar = round(size * 2.2)

	width  = cols * w + 2 * margin
	height = rows * h + 2 * margin + bar

	outside = (0x12, 0x13, 0x18)
	bar_colour = (0x2b, 0x2d, 0x3a)
	title_colour = (0xa0, 0xa4, 0xb0)

	if light:
		outside = (0xd8, 0xd8, 0xd4)
		bar_colour = (0xe6, 0xe6, 0xe2)
		title_colour = (0x60, 0x64, 0x70)

	image = Image.new("RGB", (width, height), outside)
	draw = ImageDraw.Draw(image)

	# The window
	draw.rounded_rectangle((0, 0, width - 1, height - 1), radius=size, fill=BACKGROUND)
	draw.rounded_rectangle((0, 0, width - 1, bar + size), radius=size, fill=bar_colour)
	draw.rectangle((0, bar, width - 1, bar + size), fill=BACKGROUND)

	for i, colour in enumerate(((0xff, 0x5f, 0x57), (0xfe, 0xbc, 0x2e), (0x28, 0xc8, 0x40))):
		x = margin + i * size * 1.4
		r = size * 0.45
		draw.ellipse((x - r, bar / 2 - r, x + r, bar / 2 + r), fill=colour)

	if title:
		draw.text((width / 2, bar / 2), title, font=regular, fill=title_colour, anchor="mm")

	top = bar + margin

	for y, cells in enumerate(lines):
		for x, cell in enumerate(cells):
			fg = cell.fg or FOREGROUND
			bg = cell.bg or BACKGROUND

			if cell.reverse:
				fg, bg = bg, fg

			if cell.dim:
				fg = blend(fg, bg, 0.45)

			x0 = margin + x * w
			y0 = top + y * h

			if bg != BACKGROUND:
				draw.rectangle((x0, y0, x0 + w - 1, y0 + h - 1), fill=bg)

			char = cell.char

			if char in LOWER:
				top_of_block = y0 + h - round(h * LOWER[char] / 8)
				draw.rectangle((x0, top_of_block, x0 + w - 1, y0 + h - 1), fill=fg)

			elif char in LEFT:
				draw.rectangle((x0, y0, x0 + round(w * LEFT[char] / 8) - 1, y0 + h - 1), fill=fg)

			elif char in LINES:
				draw_line(draw, char, x0, y0, w, h, fg, max(1, round(size / 12)))

			elif char.strip():
				font = bold if cell.bold else regular

				draw.text((x0, y0 + h - descent - round(size * 0.05)), char,
					font=font, fill=fg, anchor="ls")

			if cell.underline:
				draw.line((x0, y0 + h - 2, x0 + w - 1, y0 + h - 2), fill=fg)

	image.save(output, optimize=True)

def main():
	parser = argparse.ArgumentParser(description="Turns a terminal into a picture")

	parser.add_argument("input", help="What tmux capture-pane -p -e has written")
	parser.add_argument("output", help="The picture (PNG)")
	parser.add_argument("--size", type=int, default=15, help="The size of the font")
	parser.add_argument("--title", default="fireperf", help="The title of the window")
	parser.add_argument("--light", action="store_true", help="Draw a terminal with a light background")

	args = parser.parse_args()

	with open(args.input, encoding="utf-8") as f:
		lines = parse(f.read())

	render(lines, args.output, args.size, args.title, light=args.light)

main()
