#!/usr/bin/env python3
"""yazi's image and PDF previews show in a terminal pane, and go with it.

python3 -B test/yazi_preview.py <pardes-gui> <output dir> [--tty <pardes>]

A hidden PARDES_TEST window runs yazi in a terminal pane on a directory
holding a PNG and a PDF; its captures must show the picture's pixels in
the preview column, then the PDF's page, then nothing once yazi quits.
With --tty, the same yazi runs under a tty pardes nested in that pane: the
tty passes the images on as kitty graphics, which the window draws.
The captures are left in the output directory. Without yazi it skips.
"""
import argparse
import os
import shutil
import struct
import time
import zlib
from pathlib import Path

from fs import session, new_pane, execute
from ui_review import require, wait_for

RED = (255, 0, 0)
YELLOW = (255, 255, 0)
WHITE = (255, 255, 255)


def kind(r, g, b):
    """Which of the three a pixel is, near enough: scaling blurs edges."""
    if r > 230 and b < 40:
        return RED if g < 40 else YELLOW if g > 230 else None
    return WHITE if r > 245 and g > 245 and b > 245 else None


def png(width, height):
    """Red, with a yellow square in the middle."""
    def pixel(x, y):
        inside = width // 3 <= x < 2 * width // 3 and height // 3 <= y < 2 * height // 3
        return bytes(YELLOW if inside else RED)
    raw = b''.join(b'\0' + b''.join(pixel(x, y) for x in range(width)) for y in range(height))
    def chunk(kind, data):
        return struct.pack('>I', len(data)) + kind + data + struct.pack('>I', zlib.crc32(kind + data))
    return (b'\x89PNG\r\n\x1a\n' + chunk(b'IHDR', struct.pack('>IIBBBBB', width, height, 8, 2, 0, 0, 0)) +
            chunk(b'IDAT', zlib.compress(raw)) + chunk(b'IEND', b''))


def frame(root):
    """The latest capture, as (width, height, rgb bytes)."""
    data = (root / 'latest.ppm').read_bytes()
    parts = data.split(maxsplit=4)
    width, height = int(parts[1]), int(parts[2])
    pixels = parts[4]
    return width, height, pixels[:width * height * 3]


def count(root, colours):
    """How many of every third pixel of every third row are each colour."""
    try:
        width, height, pixels = frame(root)
    except (OSError, ValueError, IndexError):
        return {c: 0 for c in colours}
    found = {c: 0 for c in colours}
    for y in range(0, height, 3):
        row = pixels[y * width * 3:(y + 1) * width * 3]
        for x in range(0, len(row) - 2, 9):
            rgb = kind(row[x], row[x + 1], row[x + 2])
            if rgb in found:
                found[rgb] += 1
    return found


def save(root, label):
    shutil.copyfile(root / 'latest.ppm', root / (label + '.ppm'))
    print(label, flush=True)


def panes(client):
    rows = [row.split(maxsplit=5) for row in client.read('/index').decode().splitlines() if row.strip()]
    return {int(row[0]): (row[5] if len(row) > 5 else '') for row in rows}


def run(args):
    if shutil.which('yazi') is None:
        print('yazi is not installed: skipped')
        return
    root = args.output.resolve()
    shutil.rmtree(root, ignore_errors=True)
    root.mkdir(parents=True, mode=0o700)
    (root / 'config/pardes').mkdir(parents=True)
    (root / 'config/pardes/init').write_text('Shell /bin/sh\nTheme forge\nMessageLinger 0\n')
    pictures = root / 'pictures'
    pictures.mkdir()
    (pictures / 'a.png').write_bytes(png(240, 160))
    shutil.copyfile(Path(__file__).resolve().parent.parent / 'docs/design.pdf', pictures / 'b.pdf')
    display = Path(os.environ.get('WAYLAND_DISPLAY', 'wayland-0'))
    if not display.is_absolute():
        display = Path(os.environ.get('XDG_RUNTIME_DIR', '/run/user/%d' % os.getuid())) / display
    name = 'yazi-preview'
    inherited = {'PARDES_TEST': '1', 'PARDES_TEST_COLS': '140', 'PARDES_TEST_ROWS': '40',
                 'PARDES_TEST_CAPTURE_DIR': str(root), 'SDL_VIDEODRIVER': 'wayland',
                 'WAYLAND_DISPLAY': str(display), 'XDG_CACHE_HOME': str(root / 'cache'),
                 'XDG_DATA_HOME': str(root / 'data'), 'XDG_STATE_HOME': str(root / 'state')}
    launch = ['-c', 'stty cols 140 rows 40; exec "$@"', name, str(args.binary.resolve()), '--9p=' + name]
    with session('/bin/sh', root, name, socket_name=name, tty=True, launch=launch, inherited=inherited) as (client, _):
        scratch = new_pane(client, b'')
        before = panes(client)
        execute(client, scratch, 'Tty')
        term = require(wait_for(lambda: next((s for s in panes(client) if s not in before and
                                              'pty' in client.list(f'/pane/{s}')), None)), 'no terminal')
        client.remove(f'/pane/{scratch}')
        time.sleep(1)
        save(root, '00-before')
        type_in = lambda text: client.write(f'/pane/{term}/pty/data', text.encode())
        if args.tty:
            type_in(f'cd {pictures} && exec env -u PARDES_PID -u PARDES_9P {args.tty.resolve()} --nested\r')
            time.sleep(4)
            type_in(f'yazi\r')
        else:
            type_in(f'cd {pictures} && yazi\r')
        # The picture, red and yellow, in the preview column.
        shown = wait_for(lambda: (lambda c: c[RED] > 200 and c[YELLOW] > 50 and c)(count(root, [RED, YELLOW])), 30)
        save(root, '01-png')
        require(shown, f'the PNG preview never showed: {count(root, [RED, YELLOW])}')
        # Down to the PDF: the picture goes, a white page comes.
        type_in('j')
        page = wait_for(lambda: (lambda c: c[RED] == 0 and c[WHITE] > 80 and c)(count(root, [RED, WHITE])), 30)
        save(root, '02-pdf')
        require(page, f'the PDF preview never replaced the picture: {count(root, [RED, WHITE])}')
        # Back up, and out of yazi: nothing of either is left behind.
        type_in('k')
        require(wait_for(lambda: count(root, [RED])[RED] > 200, 30), 'the PNG did not come back')
        type_in('q')
        gone = wait_for(lambda: (lambda c: c[RED] == 0 and c[WHITE] < 20 and c)(count(root, [RED, WHITE])), 20)
        save(root, '03-quit')
        require(gone, f'images stayed after yazi quit: {count(root, [RED, WHITE])}')
    print('yazi previews showed and went' + (' through a nested tty' if args.tty else ''))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('binary', type=Path)
    parser.add_argument('output', type=Path)
    parser.add_argument('--tty', type=Path, help='a tty pardes to nest in the pane, yazi in it')
    run(parser.parse_args())