#!/usr/bin/env python3
"""
BeebWeb PC client - browse a BeebWeb server from your desktop.

It speaks the same protocol as the BBC BASIC client and mirrors that client's
behaviour tag for tag, drawing onto an emulated BBC screen (real screen memory
layouts for MODEs 0-7, hardware scrolling, VDU codes, palette, teletext), so
what you see here is what the Beeb will show.

    pip install pillow pygame
    python beebweb_pc_client.py                    # localhost:8065
    python beebweb_pc_client.py myserver 8065 --scale 2 --baud 19200

Keys (these send exactly the codes the BBC keyboard would):
    Left/Right/Down/Up   140 back / 141 links / 142 down / 143 up
    F1..F10              f0..f9  (f0 = index, f1 = home, f2 = file list)
    Shift/Ctrl + F-key   the shifted/ctrl'd soft key codes
    Ctrl+R refresh   Ctrl+L go to a page by name   Ctrl+S screenshot   Ctrl+K scanlines
    End                  the COPY key (139): pick an <input> field from a list
    Left mouse button    click a link, input field or menu item (when the pointer is on)
    Ctrl+H help      Ctrl+Q quit      Esc cancels a download
Everything else you type is sent as-is (e.g. a link number then Return).
"""
import argparse
import os
import re
import socket
import sys
import threading
import time
from collections import namedtuple

from PIL import Image, ImageChops

# ---------------------------------------------------------------------------
# 8x8 font: public domain (Marcel Sondaar / IBM VGA, via Daniel Hepper's
# font8x8). 96 glyphs for 0x20-0x7F, 8 bytes each, least significant bit =
# leftmost pixel. 0x60 is the pound sign, as on the BBC.
# ---------------------------------------------------------------------------
FONT_HEX = (
    "0000000000000000183c3c1818001800363600000000000036367f367f3636000c3e031e301f0c00006333180c666300"
    "1c361c6e3b336e000606030000000000180c0606060c1800060c1818180c060000663cff3c660000000c0c3f0c0c0000"
    "00000000000c0c060000003f0000000000000000000c0c006030180c060301003e63737b6f673e000c0e0c0c0c0c3f00"
    "1e33301c06333f001e33301c30331e00383c36337f3078003f031f3030331e001c06031f33331e003f3330180c0c0c00"
    "1e33331e33331e001e33333e30180e00000c0c00000c0c00000c0c00000c0c06180c0603060c180000003f00003f0000"
    "060c1830180c06001e3330180c000c003e637b7b7b031e000c1e33333f3333003f66663e66663f003c66030303663c00"
    "1f36666666361f007f46161e16467f007f46161e16060f003c66030373667c003333333f333333001e0c0c0c0c0c1e00"
    "7830303033331e006766361e366667000f06060646667f0063777f7f6b63630063676f7b736363001c36636363361c00"
    "3f66663e06060f001e3333333b1e38003f66663e366667001e33070e38331e003f2d0c0c0c0c1e003333333333333f00"
    "33333333331e0c006363636b7f7763006363361c1c3663003333331e0c0c1e007f6331184c667f001e06060606061e00"
    "03060c18306040001e18181818181e00081c36630000000000000000000000ff1c36260f06673f0000001e303e336e00"
    "0706063e66663b0000001e3303331e003830303e33336e0000001e333f031e001c36060f06060f0000006e33333e301f"
    "0706366e666667000c000e0c0c0c1e00300030303033331e070666361e3667000e0c0c0c0c0c1e000000337f7f6b6300"
    "00001f333333330000001e3333331e0000003b66663e060f00006e33333e307800003b6e66060f0000003e031e301f00"
    "080c3e0c0c2c18000000333333336e0000003333331e0c000000636b7f7f3600000063361c36630000003333333e301f"
    "00003f190c263f00380c0c070c0c38001818180018181800070c0c380c0c07006e3b0000000000000000000000000000"
)
FONT = [bytes.fromhex(FONT_HEX)[i * 8:(i + 1) * 8] for i in range(96)]


# ---------------------------------------------------------------------------
# The thin 4x8 font used by BeebWeb's MODE 8 / MODE 9 (80 / 40 half-width
# columns): 224 glyphs for characters 32-255, 8 bytes each, one byte per pixel
# row, the leftmost pixel in bit 7 and the four pixels in the top nibble.
# ---------------------------------------------------------------------------
THIN_HEX = (
    "00000000000000004040404040004000a0a0a000000000000040e04040e040004060804020c0400000a020404080a000"
    "40a040c0b0a0500020204000000000002040404040402000402020202020400000a040e040a00000004040e040400000"
    "0000000000202040000000e0000000000000000000004000002020404080800060a0a0a0a0a0c00040c040404040e000"
    "40a020404080e00040a0204020a04000206060a0e0202000e08080c020a0400040a080c0a0a04000e020204040808000"
    "40a0a040a0a0400040a0a06020a040000000400000004000000020000020204000204080402000000000e00000e00000"
    "008040204080000040a020404000400040a0a0e0e080600040a0a0a0e0a0a000c0a0a0c0a0a0c00040a0808080a04000"
    "c0a0a0a0a0a0c000e08080c08080e000e08080c08080800040a08080a0a06000a0a0a0e0a0a0a000e04040404040e000"
    "e040404040408000a0a0a0c0a0a0a000808080808080e000a0e0e0a0a0a0a000c0a0a0a0a0a0a00040a0a0a0a0a04000"
    "c0a0a0c08080800040a0a0a0e0c06000c0a0a0c0a0a0a00040a0804020a04000e040404040404000a0a0a0a0a0a04000"
    "a0a0a0a0a0404000a0a0a0a0e0e0a000a0a0404040a0a000a0a0a0a040404000e02020404080e000e08080808080e000"
    "0080804040202000e02020202020e00040a000000000000000000000000000f0204040e04040e0000000c02060a06000"
    "8080c0a0a0a0c0000000608080806000202060a0a0a06000000040a0e0806000204040e040404000000060a0a06020c0"
    "8080c0a0a0a0a0000040004040404000002000202020a040808080a0c0a0a00040404040404020000000a0e0e0a0a000"
    "0000c0a0a0a0a000000040a0a0a040000000c0a0a0c08080000060a0a06030200000a0c080808000000060804020c000"
    "0040e040404020000000a0a0a0a060000000a0a0a04040000000a0a0e0e0a0000000a0a040a0a0000000a0a0a06020c0"
    "0000e0204080e00020404080404020004040400040404000804040204040800050a0000000000000f0f0f0f0f0f0f0f0"
    "a00040a0e0a0a00040a040a0e0a0a00050a0a0b0e0a0b00040a08080a04040802040e080c080e000a040a0a0a0a04000"
    "a000a0a0a0a040006090b0d0b0906000000040f040000000000020f0200000004040404040e040000040e04040404040"
    "4020c02060a06000402040a0e0806000a00040a0e080600040a040a0e0806000a000c02060a0600040a0c02060a06000"
    "0000f05070a070000000608080604080204040a0e0806000a00040a0a0a04000a000a0a0a0a060008040004040404000"
    "40a0004040404000402040a0a0a0400040a040a0a0a04000402000a0a0a0600040a000a0a0a06000a000a0a0a06020c0"
    "0090609090609000608060a0c020c0006090600000000000000000606060606000000070000000000000007060606060"
    "000000e000000000000000e060606060000000f000000000000000f06060606060606060000000006060606060606060"
    "60606070000000006060607060606060606060e000000000606060e060606060606060f000000000606060f060606060"
    "000000103060606000000080c060606060603010000000006060c08000000000004000404080a0400040004040404040"
    "50a000c0a0a0a00050a00040a0a0a00060f060606060600060f0606060f0600060606000000000004040200000000000"
    "70d0d0d0505070000000006060000000101020a040400000a050a050a050a050c0a0a0a0a0a0600040a0a0a0e0a0a000"
    "c0a0a0c0a0a0c000e080808080808000404040a0a0a0e000e08080c08080e000e02020404080e000a0a0a0e0a0a0a000"
    "40a0a0e0a0a04000e04040404040e000a0a0a0c0a0a0a000404040a0a0a0a000a0e0e0a0a0a0a000c0a0a0a0a0a0a000"
    "e00000400000e00040a0a0a0a0a04000f0a0a0a0a0a0a000c0a0a0c080808000e08040204080e000e040404040404000"
    "a0a0a0a040404000e040e0a0e040e000a0a0404040a0a000e04040a0e040e00040a0a0a0a040a000e0a0a0a040404000"
    "0040e04000e0000000e00040e04000006060606060606000a0a0a0a0a0a0a0000000a0a0a0a04000000040a0a0a0a000"
    "2040a0a0a040800000000050a0a050000000c0a0e0a0c0800000a0a04040808000608040a0a0400000006080e0806000"
    "00e02040808060200000c0a0a0a0a020000040a0e0a040000000004040404000000000a0c0a0a0000080804040a0a000"
    "0000a0a0a0a0c0800000a0a06060200000006080e0806020000040a0a0a04000000000f0a0a0a000000040a0a0c08080"
    "000070a0a0a040000000e040404020000000a0a0a0a060000040a0a0404040000000a0a040a0a000004040a0e0404000"
    "0000a0a0e0e0a00000c02040a0a04000000050a000f0000000e000e000e00000204080402000e0000020e04040e08000"
    "804020408000e000f0f0f0f0f0f0f0f0"
)
THIN_CHARS = range(32, 256)
THIN = [bytes.fromhex(THIN_HEX)[i * 8:(i + 1) * 8] for i in range(224)]

PHYS_RGB = [(0, 0, 0), (255, 0, 0), (0, 255, 0), (255, 255, 0),
            (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255)]
DEFAULT_PAL = {2: [0, 7], 4: [0, 1, 3, 7],
               16: [0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7]}

Mode = namedtuple("Mode", "base size cols rows bcols ppb ncol pxw scan cellw")
#                          base    size   cols rows bcols ppb ncol pxw  scan cellw
MODES = {
    0: Mode(0x3000, 0x5000, 80, 32, 80, 8, 2, 640, 8, 8),
    1: Mode(0x3000, 0x5000, 40, 32, 80, 4, 4, 320, 8, 8),
    2: Mode(0x3000, 0x5000, 20, 32, 80, 2, 16, 160, 8, 8),
    3: Mode(0x4000, 0x4000, 80, 25, 80, 8, 2, 640, 10, 8),
    4: Mode(0x5800, 0x2800, 40, 32, 40, 8, 2, 320, 8, 8),
    5: Mode(0x5800, 0x2800, 20, 32, 40, 4, 4, 160, 8, 8),
    6: Mode(0x6000, 0x2000, 40, 25, 40, 8, 2, 320, 10, 8),
    7: Mode(0x7C00, 0x0400, 40, 25, 40, 0, 8, 0, 0, 8),   # teletext
    # BeebWeb's extra modes: the pixel layout and colours of MODE 1 / MODE 2, but with
    # twice as many (half-width, 4 pixel) text columns, i.e. 80 and 40 columns.
    8: Mode(0x3000, 0x5000, 80, 32, 80, 4, 4, 320, 8, 4),
    9: Mode(0x3000, 0x5000, 40, 32, 80, 2, 16, 160, 8, 4),
}

# --- pixel <-> byte tables, using the real BBC screen byte layouts ---------
def _unpack(ppb, b):
    bit = lambda n: (b >> n) & 1
    if ppb == 8:
        return bytes(bit(7 - n) for n in range(8))
    if ppb == 4:
        return bytes((bit(7 - n) << 1) | bit(3 - n) for n in range(4))
    return bytes(((bit(7 - n) << 3) | (bit(5 - n) << 2) |
                  (bit(3 - n) << 1) | bit(1 - n)) for n in range(2))

UNPACK = {p: [_unpack(p, b) for b in range(256)] for p in (8, 4, 2)}
PACK = {p: {tuple(UNPACK[p][b]): b for b in range(256)} for p in (8, 4, 2)}


def bbc_val(s):
    """BBC BASIC VAL(): leading number, &hex allowed, else 0."""
    m = re.match(r"\s*([-+]?(?:\d+\.?\d*|\.\d+)(?:[eE][-+]?\d+)?)", s)
    if m:
        return float(m.group(1))
    m = re.match(r"\s*&([0-9A-Fa-f]+)", s)
    return int(m.group(1), 16) if m else 0


class Disconnected(Exception):
    pass


# ---------------------------------------------------------------------------
# The emulated BBC screen
# ---------------------------------------------------------------------------
VDU_PARAMS = {17: 1, 18: 2, 19: 5, 22: 1, 23: 9, 24: 8, 25: 5, 28: 4, 29: 4, 31: 2}


class Screen:
    def __init__(self):
        self.mem = bytearray(0x10000)
        self.lock = threading.RLock()
        self.dirty = True
        self.udg = {}
        self._cache = {}
        self._pend = None
        self._masks = {}
        self.scanlines = 0.4        # 0 = off; darkness of every second output row
        self._slmask = {}
        self.set_mode(1)

    # -- addressing ---------------------------------------------------------
    def _idx(self, off):
        m = self.info
        return (self.start - m.base + off) % m.size

    def addr(self, off):
        return self.info.base + self._idx(off)

    def _fill(self, off, n, val):
        m = self.info
        i = self._idx(off)
        first = min(n, m.size - i)
        self.mem[m.base + i:m.base + i + first] = bytes([val]) * first
        if n > first:
            self.mem[m.base:m.base + n - first] = bytes([val]) * (n - first)

    @property
    def stride(self):
        return 40 if self.mode == 7 else self.info.bcols * 8

    @property
    def bgbyte(self):
        if self.mode == 7:
            return 0x20
        return PACK[self.info.ppb][(self.bg,) * self.info.ppb]

    # -- modes / colours ----------------------------------------------------
    def set_mode(self, n):
        with self.lock:
            self.mode_req = n
            if n not in MODES:
                n = n & 7
            self.mode = n
            self.info = MODES[n]
            self.pal = list(DEFAULT_PAL.get(self.info.ncol, DEFAULT_PAL[2]))
            self.fg = self.info.ncol - 1
            self.bg = 0
            self.cls()

    def colour(self, n):
        n = int(n)
        if self.mode == 7:
            return
        mask = self.info.ncol - 1
        if n < 128:
            self.fg = n & mask
        else:
            self.bg = (n - 128) & mask

    def cls(self):
        with self.lock:
            m = self.info
            self.start = m.base
            self.mem[m.base:m.base + m.size] = bytes([self.bgbyte]) * m.size
            self.cx = self.cy = 0
            self.dirty = True

    # -- VDU stream ---------------------------------------------------------
    def put(self, b):
        with self.lock:
            if self._pend is not None:
                self._pend.append(b)
                if len(self._pend) - 1 == VDU_PARAMS[self._pend[0]]:
                    code, args = self._pend[0], self._pend[1:]
                    self._pend = None
                    self._vdu(code, args)
            elif b < 32:
                if b in VDU_PARAMS:
                    self._pend = [b]
                else:
                    self._vdu(b, [])
            else:
                self._print(b)
            self.dirty = True

    def print_str(self, s, newline=True):
        """BASIC PRINT: text goes through the VDU drivers; newline = CR LF."""
        for ch in s:
            self.put(ord(ch) & 255)
        if newline:
            self.put(10)
            self.put(13)

    def _vdu(self, code, a):
        m = self.info
        if code == 8:
            if self.cx > 0:
                self.cx -= 1
            elif self.cy > 0:
                self.cx, self.cy = m.cols - 1, self.cy - 1
        elif code == 9:
            self._advance()
        elif code == 10:
            self._down()
        elif code == 11:
            self.cy = max(0, self.cy - 1)
        elif code == 12:
            self.cls()
        elif code == 13:
            self.cx = 0
        elif code == 17:
            self.colour(a[0])
        elif code == 19:
            if self.mode != 7:
                self.pal[a[0] & (m.ncol - 1)] = a[1] & 7
        elif code == 20:
            self.fg, self.bg = m.ncol - 1, 0
            self.pal = list(DEFAULT_PAL.get(m.ncol, DEFAULT_PAL[2]))
        elif code == 22:
            self.set_mode(a[0])
        elif code == 23:
            if a[0] >= 32:
                # BBC rows have bit 7 = leftmost; our font has bit 0 = leftmost
                self.udg[a[0]] = bytes(int(f"{r:08b}"[::-1], 2) for r in a[1:9])
                self._cache.clear()
                self._masks.clear()
        elif code == 30:
            self.cx = self.cy = 0
        elif code == 31:
            self.cx = min(a[0], m.cols - 1)
            self.cy = min(a[1], m.rows - 1)

    def _down(self):
        self.cy += 1
        if self.cy >= self.info.rows:
            self.cy = self.info.rows - 1
            self.scroll()

    def _advance(self):
        self.cx += 1
        if self.cx >= self.info.cols:
            self.cx = 0
            self._down()

    def scroll(self):
        """Hardware scroll: move the display start address on a line, clear the new last line."""
        m = self.info
        self.start = m.base + (self.start - m.base + self.stride) % m.size
        self._fill((m.rows - 1) * self.stride, self.stride, self.bgbyte)

    def _print(self, c):
        if self.mode == 7:
            self.mem[self.addr(self.cy * 40 + self.cx)] = c
        else:
            self._plot(c, self.cx, self.cy)
        self._advance()

    def glyph(self, c):
        if c in self.udg:
            return self.udg[c]
        if 32 <= c <= 127:
            return FONT[c - 32]
        return bytes(8)

    def _plot(self, c, col, row):
        m = self.info
        key = (c, self.fg, self.bg, m.ppb, m.cellw)
        cells = self._cache.get(key)
        per = m.cellw // m.ppb
        if cells is None:
            g = self.glyph(c)
            cells = []
            for s in range(8):
                if m.cellw == 8:
                    on = [(g[s] >> x) & 1 for x in range(8)]
                elif c in THIN_CHARS and c not in self.udg:   # half-width cell: the real thin font
                    on = [(THIN[c - 32][s] >> (7 - x)) & 1 for x in range(4)]
                else:   # a user defined character in a half-width cell: sample every
                        # other column; a row that would vanish ORs pairs instead
                    on = [(g[s] >> (2 * x)) & 1 for x in range(4)]
                    if g[s] and not any(on):
                        on = [((g[s] >> (2 * x)) | (g[s] >> (2 * x + 1))) & 1 for x in range(4)]
                px = [self.fg if v else self.bg for v in on]
                cells.append([PACK[m.ppb][tuple(px[k * m.ppb:(k + 1) * m.ppb])]
                              for k in range(per)])
            self._cache[key] = cells
        base = row * self.stride + col * per * 8
        for s in range(8):
            for k in range(per):
                self.mem[self.addr(base + k * 8 + s)] = cells[s][k]

    def poke_block(self, addr, data):
        """Raw memory write (<IMGSTART>): ROM space above &8000 ignores writes."""
        with self.lock:
            data = bytes(data)[:max(0, 0x8000 - addr)]
            self.mem[addr:addr + len(data)] = data
            self.dirty = True

    def cell_at(self, x, y):
        """Text cell under a point of the 640x512 picture, or None."""
        m = self.info
        if self.mode == 7:
            col, row = x // 16, (y - 6) // 20
        else:
            h = m.rows * m.scan
            col = x * m.cols // 640
            row = (y - (512 - h * 2) // 2) // (m.scan * 2)
        if 0 <= col < m.cols and 0 <= row < m.rows:
            return int(col), int(row)
        return None

    # -- rendering ----------------------------------------------------------
    def render(self):
        img = self._render_raw()
        d = self.scanlines
        if d > 0:
            mk = self._slmask.get(d)
            if mk is None:
                dark = int(255 * (1 - d))
                row_a, row_b = bytes([255]) * 640, bytes([dark]) * 640
                mk = Image.frombytes("L", (640, 512), (row_a + row_b) * 256).convert("RGB")
                self._slmask[d] = mk
            img = ImageChops.multiply(img, mk)
        return img

    def _render_raw(self):
        with self.lock:
            m = self.info
            snap = (self.mode, m, bytes(self.mem[m.base:m.base + m.size]),
                    (self.start - m.base) % m.size, list(self.pal))
            self.dirty = False
        mode, m, buf, st, pal = snap

        def region(off, n):
            i = (st + off) % m.size
            if i + n <= m.size:
                return buf[i:i + n]
            return buf[i:] + buf[:i + n - m.size]

        if mode == 7:
            return self._render_teletext(region)
        table = UNPACK[m.ppb]
        stride = m.bcols * 8
        out = bytearray()
        gap = bytes([16]) * m.pxw
        for r in range(m.rows):
            rb = region(r * stride, stride)
            for s in range(8):
                out += b"".join([table[v] for v in rb[s::8]])
            for _ in range(m.scan - 8):
                out += gap
        h = m.rows * m.scan
        img = Image.frombytes("P", (m.pxw, h), bytes(out))
        flat = []
        for i in range(m.ncol):
            flat += PHYS_RGB[pal[i] & 7]
        flat += [0, 0, 0] * (17 - m.ncol)
        img.putpalette(flat)
        img = img.convert("RGB").resize((640, h * 2), Image.NEAREST)
        canvas = Image.new("RGB", (640, 512))
        canvas.paste(img, (0, (512 - h * 2) // 2))
        return canvas

    # teletext ---------------------------------------------------------------
    def _glyph_mask(self, c):
        mk = self._masks.get(("g", c))
        if mk is None:
            mk = Image.new("L", (16, 20), 0)
            g = self.glyph(c)
            for y in range(8):
                for x in range(8):
                    if (g[y] >> x) & 1:
                        mk.paste(255, (2 * x, 2 * y + 2, 2 * x + 2, 2 * y + 4))
            self._masks[("g", c)] = mk
        return mk

    def _mosaic_mask(self, c, sep):
        key = ("m", c, sep)
        mk = self._masks.get(key)
        if mk is None:
            mk = Image.new("L", (16, 20), 0)
            bits = (c & 0x1F) | ((c >> 6) & 1) << 5
            xs = [(0, 8), (8, 16)]
            ys = [(0, 7), (7, 13), (13, 20)]
            for i in range(6):
                if bits >> i & 1:
                    x0, x1 = xs[i % 2]
                    y0, y1 = ys[i // 2]
                    if sep:
                        x1 -= 2
                        y1 -= 2
                    mk.paste(255, (x0, y0, x1, y1))
            self._masks[key] = mk
        return mk

    def _half_mask(self, mask, half):
        key = ("h", id(mask), half)
        mk = self._masks.get(key)
        if mk is None:
            box = (0, 0, 16, 10) if half == "top" else (0, 10, 16, 20)
            mk = mask.crop(box).resize((16, 20), Image.NEAREST)
            self._masks[key] = mk
        return mk

    def _render_teletext(self, region):
        img = Image.new("RGB", (640, 512))
        data = region(0, 1000)

        def row_render(rb, y, half):
            fg, bg, gfx, sep, dh, hold, held, held_sep, conc = 7, 0, False, False, False, False, 0x20, False, False
            for c in range(40):
                b = rb[c] & 0x7F
                ctrl = b < 32
                if ctrl:                       # "set-at" codes
                    if b == 25:
                        sep = False
                    elif b == 26:
                        sep = True
                    elif b == 28:
                        bg = 0
                    elif b == 29:
                        bg = fg
                    elif b == 30:
                        hold = True
                show, ssep = None, sep
                if ctrl:
                    if hold and gfx:
                        show, ssep = held, held_sep
                else:
                    show = b
                    if gfx and (b < 0x40 or b >= 0x60):
                        held, held_sep = b, sep
                x = c * 16
                box = (x, y, x + 16, y + 20)
                if bg:
                    img.paste(PHYS_RGB[bg], box)
                if show is not None and show != 0x20 and not conc and (dh or half != "bottom"):
                    if gfx and (0x20 <= show < 0x40 or 0x60 <= show < 0x80):
                        mk = self._mosaic_mask(show, ssep)
                    else:
                        mk = self._glyph_mask(show)
                    if half and dh:
                        mk = self._half_mask(mk, half)
                    img.paste(PHYS_RGB[fg], box, mk)
                if ctrl:                       # "set-after" codes
                    if 1 <= b <= 7:
                        fg, gfx = b, False
                    elif 17 <= b <= 23:
                        fg, gfx = b - 16, True
                    elif b == 12:
                        dh = False
                    elif b == 13:
                        dh = True
                    elif b == 24:
                        conc = True
                    elif b == 31:
                        hold = False

        prev = None
        for r in range(25):
            rb = data[r * 40:(r + 1) * 40]
            y = 6 + r * 20
            if prev is not None:
                row_render(prev, y, "bottom")
                prev = None
            elif any((v & 0x7F) == 13 for v in rb):
                row_render(rb, y, "top")
                prev = rb
            else:
                row_render(rb, y, None)
        return img


# ---------------------------------------------------------------------------
# The protocol client: a faithful port of the BBC BASIC client
# ---------------------------------------------------------------------------
ACK, NAK, CAN = 6, 21, 24
BLOCK = 128
BYTE_TIMEOUT = 5.0
MAX_TRIES = 10


class Client:
    """
    ui must provide:  ask_save(desc, size) -> path or None   (may block)
                      cancel_requested() -> bool
    """

    def __init__(self, host, port, screen, ui, baud=0):
        self.host, self.port = host, port
        self.screen = screen
        self.ui = ui
        self.baud = baud
        self.sock = None
        self.buf = bytearray()
        self.pos = 0
        self._due = 0.0
        self.alive = False
        self.events = []            # (time, text) - shown in the status bar
        self.last_rx = time.monotonic()
        self.transfer = None        # text while a download is in progress
        self.mouse_on = False       # the pointer is live (<pb> / <end>); text or <moff> hides it
        self.attr = None            # (load, exec) from <attr>, for the next download
        # BASIC variables
        self.M, self.D, self.B, self.F = 0, 1, 128, 3

    def note(self, text):
        self.events.append((time.time(), text))
        del self.events[:-50]

    # -- connection ---------------------------------------------------------
    def connect(self):
        self.sock = socket.create_connection((self.host, self.port), timeout=10)
        self.sock.settimeout(None)
        self.alive = True
        threading.Thread(target=self._run, daemon=True).start()

    def close(self):
        self.alive = False
        try:
            self.sock.close()
        except Exception:
            pass

    def send(self, data):
        if isinstance(data, str):
            data = data.encode("latin-1", "replace")
        try:
            self.sock.sendall(data)
        except OSError:
            self.alive = False

    # -- byte input -----------------------------------------------------------
    def _fill(self, timeout):
        self.sock.settimeout(timeout)
        try:
            d = self.sock.recv(4096)
        except socket.timeout:
            return False
        except OSError:
            raise Disconnected
        if not d:
            raise Disconnected
        self.buf = bytearray(d)
        self.pos = 0
        return True

    def _throttle(self, k):
        if self.baud:
            now = time.monotonic()
            self._due = max(self._due, now) + k * 10.0 / self.baud
            delay = self._due - now
            if delay > 0.003:
                time.sleep(delay)

    def read_byte(self, timeout=None):
        if self.pos >= len(self.buf) and not self._fill(timeout):
            return None
        b = self.buf[self.pos]
        self.pos += 1
        self.last_rx = time.monotonic()
        self._throttle(1)
        return b

    def read_exact(self, n):
        out = bytearray()
        while len(out) < n:
            if self.pos >= len(self.buf):
                self._fill(None)
            take = min(n - len(out), len(self.buf) - self.pos)
            out += self.buf[self.pos:self.pos + take]
            self.pos += take
            self._throttle(take)
        self.last_rx = time.monotonic()
        return bytes(out)

    def flush_input(self, quiet=0.0):
        """*FX21 - throw away anything waiting (optionally until the line is quiet)."""
        self.buf, self.pos = bytearray(), 0
        while self._fill(quiet if quiet else 0.001):
            self.buf, self.pos = bytearray(), 0

    # -- main loop ------------------------------------------------------------
    def _run(self):
        try:
            while self.alive:
                b = self.read_byte()
                if b == 60:                         # '<'
                    self.tag()
                elif b == 13:                       # OSASCI turns CR into LF CR
                    self.mouse_on = False
                    self.screen.put(10)
                    self.screen.put(13)
                else:
                    self.mouse_on = False
                    self.screen.put(b)
        except Disconnected:
            self.note("Disconnected from server")
        except Exception as e:                      # keep the window alive for inspection
            self.note(f"Client error: {e!r}")
            import traceback
            traceback.print_exc()
        self.alive = False

    def tag(self):
        chars = []
        while True:
            b = self.read_byte()
            if b == 62:                             # '>'
                break
            chars.append(b)
        t = bytes(chars).decode("latin-1")
        sp = t.find(" ")
        name, rest = (t[:sp], t[sp + 1:]) if sp >= 0 else (t, t)
        s = self.screen
        if name in ("new", "break"):
            s.cls()
        elif name == "pb":
            self.p_pb()
            self.mouse_on = True
        elif name == "end":
            self.mouse_on = True
        elif name == "moff":
            self.mouse_on = False
        elif name == "attr":
            a, _, e = rest.partition(" ")
            try:
                self.attr = (int(a.strip().lstrip("&"), 16), int(e.strip().lstrip("&"), 16))
            except ValueError:
                self.attr = None
        elif name == "PRNDATA":
            self.p_prn(rest)
        elif name == "goodbye":
            self.note("Server said goodbye")
            self.alive = False
            raise Disconnected
        elif name == "pos":
            h, _, v = rest.partition(" ")
            with s.lock:
                s.put(31); s.put(int(bbc_val(h)) & 255); s.put(int(bbc_val(v)) & 255)
        elif name == "mode":
            s.set_mode(int(bbc_val(rest))); self.M = int(bbc_val(rest))
        elif name == "pal":
            # (no newline: the BBC client's FNpal only does the VDU 19s, and the server
            # counts a <pal> line as taking no row)
            for i, p in enumerate(rest.split(",")):
                s.put(19); s.put(i); s.put(int(bbc_val(p)) & 255)
                for _ in range(3):
                    s.put(0)
        elif name == "colour":
            s.colour(int(bbc_val(rest))); self.D = int(bbc_val(rest))
        elif name == "link":
            self.p_link(rest)
        elif name == "linkbg":
            self.B = int(bbc_val(rest)) + 128
        elif name == "linkfg":
            self.F = int(bbc_val(rest))
        elif name == "IMGSTART":
            self.p_img(rest)
        elif name == "title":
            s.put(31); s.put(0); s.put(0)
        elif name == "FILEPROMPT":
            self.p_file(rest)
            self.attr = None
        else:
            self.note(f"Ignored unknown tag <{t}>")

    def p_pb(self):
        s, D = self.screen, self.D
        if self.M == 0:
            s.colour(0); s.colour(129)
        else:
            s.colour(D + 129)
        s.colour(D); s.colour(128)

    def p_link(self, rest):
        s = self.screen
        sp = rest.find(" ")
        desc = rest[sp + 1:] if sp >= 0 else rest
        s.colour(2)
        if self.M == 0:
            s.colour(129)
        else:
            s.colour(self.B); s.colour(self.F)
        s.print_str(desc, newline=False)
        s.colour(128); s.colour(self.D)

    @staticmethod
    def _addr(txt):
        txt = txt.strip()
        if txt.startswith("&"):
            return int(txt[1:], 16)
        return int(bbc_val(txt))

    def p_img(self, rest):
        a, _, e = rest.partition(" ")
        start, end = self._addr(a), self._addr(e)
        n = max(end - start, 1)
        self.screen.poke_block(start, self.read_exact(n))
        chars = []
        while True:                                 # FNp: read up to '>'
            b = self.read_byte()
            if b == 62:
                break
            chars.append(b)
        row = int(bbc_val(bytes(chars).decode("latin-1").split(" ")[-1]))
        with self.screen.lock:
            self.screen.put(31); self.screen.put(0); self.screen.put(row & 255)

    # -- file download (port of PROCf) ----------------------------------------
    def p_file(self, rest):
        s = self.screen
        sp = rest.find(" ")
        size = int(bbc_val(rest[:sp])) if sp >= 0 else 0
        desc = rest[sp + 1:] if sp >= 0 else ""
        s.print_str("")
        s.print_str(f"Download: {desc}")
        s.print_str(f"{size} bytes")
        path = self.ui.ask_save(desc, size)
        if not path:
            self.send("<CANCEL>\r")
            s.print_str("Cancelled.")
            return
        try:
            d = os.path.dirname(os.path.abspath(path))
            os.makedirs(d, exist_ok=True)
            f = open(path, "wb")
        except OSError as e:
            self.send("<CANCEL>\r")
            s.print_str(f"Error: {e.strerror or e}")
            self.note(f"Cannot open {path}: {e}")
            return
        if size == 0:
            f.close()
            s.print_str("Empty file saved")
            return
        self.send("<READY>\r")
        s.print_str("Receiving... (ESCAPE cancels)")
        ok = self._receive(f, size)
        if ok:
            s.print_str("")
            s.print_str(f"Saved {size} bytes to {path}")
            self.note(f"Saved {size} bytes to {path}")
            if self.attr:                           # BBC load / exec addresses
                try:
                    with open(path + ".inf", "w") as inf:
                        inf.write(f"{os.path.basename(path)} {self.attr[0]:08X} {self.attr[1]:08X}\n")
                    s.print_str(f"Load &{self.attr[0]:X} exec &{self.attr[1]:X} (.inf saved)")
                except OSError:
                    pass

    def _receive(self, f, size):
        """The block receive loop shared by downloads and print jobs. Closes f."""
        s = self.screen
        left, exp, tries = size, 0, 0
        ok = False
        try:
            while left > 0:
                if self.ui.cancel_requested():
                    raise RuntimeError("Cancelled")
                n = min(BLOCK, left)
                status, seq, data = self._read_frame(n)
                if status == 0 and seq == exp:
                    f.write(data)
                    exp = (exp + 1) & 255
                    left -= n
                    tries = 0
                    self.send(bytes([ACK]))
                    self.transfer = f"{size - left}/{size}"
                elif status == 0 and seq == (exp - 1) & 255:
                    self.send(bytes([ACK]))         # duplicate: our ACK was lost
                else:
                    tries += 1
                    self.note(f"Block error ({'timeout' if status == 1 else 'checksum/sequence'}), retry {tries}")
                    if tries > MAX_TRIES:
                        raise RuntimeError("Too many errors")
                    self.flush_input(0.3 if status == 1 else 0)
                    self.send(bytes([NAK]))
            ok = True
        except Disconnected:
            raise
        except Exception as e:
            self.send(bytes([CAN]))
            s.print_str("")
            s.print_str(f"Error: {e}")
            self.note(f"Transfer failed: {e}")
        finally:
            f.close()
            self.transfer = None
        return ok

    def p_prn(self, rest):
        """<PRNDATA size>: a print job (Epson bytes). There is no printer here, so the
        job is saved as a .prn file in the downloads folder."""
        s = self.screen
        size = int(bbc_val(rest))
        d = os.path.abspath(self.ui.downloads)
        path = os.path.join(d, time.strftime("print-%Y%m%d-%H%M%S.prn"))
        try:
            os.makedirs(d, exist_ok=True)
            f = open(path, "wb")
        except OSError as e:
            self.send("<CANCEL>\r")
            self.note(f"Cannot save print job: {e}")
            return
        if size == 0:
            f.close()
            return
        self.send("<READY>\r")
        s.print_str("")
        s.print_str("Receiving print job... (ESCAPE cancels)")
        if self._receive(f, size):
            s.print_str(f"Print job saved to {path}")
            self.note(f"Print job ({size} bytes) saved to {path}")

    def click(self, col, row):
        """A left click on text cell (col, row), if the pointer is live."""
        if self.alive and self.mouse_on:
            self.mouse_on = False
            self.send(f"MSE {col} {row}\r")
            return True
        return False

    def _read_frame(self, n):
        """-> (status 0 ok / 1 timeout / 2 bad checksum, seq, data)"""
        seq = self.read_byte(BYTE_TIMEOUT)
        if seq is None:
            return 1, 0, b""
        data = bytearray()
        while len(data) < n + 1:
            b = self.read_byte(BYTE_TIMEOUT)
            if b is None:
                return 1, seq, b""
            data.append(b)
        chk = data.pop()
        if (seq + sum(data)) & 255 != chk:
            return 2, seq, b""
        return 0, seq, bytes(data)


# ---------------------------------------------------------------------------
# Front ends
# ---------------------------------------------------------------------------
def safe_name(desc):
    n = re.sub(r"[^\w.\-]+", "_", desc).strip("_")[:40]
    return n or "download.bin"


HELP = [
    "Left  = back (140)          Right = links (141)",
    "Down  = next screen (142)   Up    = previous screen (143)",
    "End = COPY key: pick an <input> field    Mouse = click links / fields",
    "F1 = f0 index   F2 = f1 home   F3 = f2 file list   F4..F10 = f3..f9",
    "Shift/Ctrl + F-key = shifted / ctrl soft keys",
    "Ctrl+R refresh   Ctrl+L go to page   Ctrl+S screenshot   Ctrl+K scanlines",
    "Esc cancels a download      Ctrl+Q quit",
    "Anything else is sent as typed - e.g. a link number, then Return.",
]


class GuiUI:
    def __init__(self, downloads):
        self.downloads = downloads
        self.entry = None           # active text prompt
        self.cancel = False

    def cancel_requested(self):
        c, self.cancel = self.cancel, False
        return c

    def ask_save(self, desc, size):
        default = os.path.join(self.downloads, safe_name(desc))
        ev = threading.Event()
        self.entry = {"label": f"Save {size} bytes as (blank = cancel):",
                      "text": default, "event": ev, "result": None, "kind": "save"}
        ev.wait()
        r = self.entry["result"]
        self.entry = None
        return r or None


def run_gui(args):
    import pygame
    pygame.init()
    pygame.display.set_caption("BeebWeb")
    scale = args.scale
    W, H = int(640 * scale), int(512 * scale)
    BAR = 46
    win = pygame.display.set_mode((W, H + BAR))
    font = pygame.font.Font(None, 22)
    small = pygame.font.Font(None, 18)

    screen = Screen()
    screen.scanlines = args.scanlines
    ui = GuiUI(args.downloads)
    cl = Client(args.host, args.port, screen, ui, args.baud)
    try:
        cl.connect()
    except OSError as e:
        print(f"Could not connect to {args.host}:{args.port}: {e}")
        return 1

    typed = ""
    show_help = False
    hand = False
    last_img = None
    clock = pygame.time.Clock()
    shots = os.path.join(os.getcwd(), "screenshots")

    def screenshot():
        os.makedirs(shots, exist_ok=True)
        p = os.path.join(shots, time.strftime("beeb-%Y%m%d-%H%M%S.png"))
        screen.render().save(p)
        cl.note(f"Screenshot saved: {p}")

    def go_page():
        ev = threading.Event()
        ui.entry = {"label": "Go to page:", "text": "", "event": ev, "result": None, "kind": "go"}

    running = True
    while running:
        for ev in pygame.event.get():
            if ev.type == pygame.QUIT:
                running = False
            elif ev.type == pygame.MOUSEBUTTONDOWN and ev.button == 1:
                if not ui.entry and not show_help and ev.pos[1] < H:
                    cell = screen.cell_at(int(ev.pos[0] / scale), int(ev.pos[1] / scale))
                    if cell and cl.click(*cell):
                        typed = ""
            elif ev.type == pygame.KEYDOWN:
                mods = pygame.key.get_mods()
                ctrl, shift = mods & pygame.KMOD_CTRL, mods & pygame.KMOD_SHIFT
                e = ui.entry
                if e:                                   # text prompt has the keyboard
                    if ev.key == pygame.K_RETURN:
                        e["result"] = e["text"].strip()
                        if e["kind"] == "go":
                            ui.entry = None
                            if e["result"]:
                                cl.send(e["result"] + "\r")
                        else:
                            e["event"].set()
                    elif ev.key == pygame.K_ESCAPE:
                        e["result"] = None
                        if e["kind"] == "go":
                            ui.entry = None
                        else:
                            e["event"].set()
                    elif ev.key == pygame.K_BACKSPACE:
                        e["text"] = e["text"][:-1]
                    elif ev.unicode and ev.unicode.isprintable():
                        e["text"] += ev.unicode
                    continue
                if show_help:
                    show_help = False
                    continue
                code = None
                if ctrl:
                    k = ev.key
                    if k == pygame.K_q:
                        running = False
                    elif k == pygame.K_s:
                        screenshot()
                    elif k == pygame.K_l:
                        go_page()
                    elif k == pygame.K_r:
                        code = 149
                    elif k == pygame.K_h:
                        show_help = True
                    elif k == pygame.K_k:
                        screen.scanlines = 0 if screen.scanlines else (args.scanlines or 0.4)
                        screen.dirty = True
                if pygame.K_F1 <= ev.key <= pygame.K_F10:
                    code = 128 + ev.key - pygame.K_F1 + (16 if shift else 0) + (32 if ctrl else 0)
                elif code is None and not ctrl:
                    code = {pygame.K_LEFT: 140, pygame.K_RIGHT: 141, pygame.K_DOWN: 142,
                            pygame.K_UP: 143, pygame.K_PAGEDOWN: 142, pygame.K_PAGEUP: 143,
                            pygame.K_RETURN: 13, pygame.K_KP_ENTER: 13,
                            pygame.K_BACKSPACE: 127, pygame.K_DELETE: 127,
                            pygame.K_END: 139}.get(ev.key)
                    if ev.key == pygame.K_ESCAPE:
                        if cl.transfer:
                            ui.cancel = True
                        else:
                            code = 27
                    elif code is None and ev.unicode and 32 <= ord(ev.unicode[0]) < 127:
                        code = ord(ev.unicode[0])
                if code is not None and cl.alive:
                    cl.send(bytes([code]))
                    if code == 13:
                        typed = ""
                    elif code == 127:
                        typed = typed[:-1]
                    elif 32 <= code < 127:
                        typed += chr(code)
                    else:
                        typed = ""

        want = bool(cl.mouse_on)                       # a hand over links while the pointer is live
        if want != hand:
            hand = want
            try:
                pygame.mouse.set_cursor(pygame.SYSTEM_CURSOR_HAND if hand else pygame.SYSTEM_CURSOR_ARROW)
            except Exception:
                pass
        if screen.dirty or last_img is None:
            img = screen.render()
            surf = pygame.image.frombuffer(img.tobytes(), img.size, "RGB")
            last_img = pygame.transform.scale(surf, (W, H)) if scale != 1 else surf
        win.fill((28, 28, 32))
        win.blit(last_img, (0, 0))

        # status bar
        pygame.draw.rect(win, (28, 28, 32), (0, H, W, BAR))
        status = "connected" if cl.alive else "DISCONNECTED"
        msg = ""
        if cl.events and time.time() - cl.events[-1][0] < 8:
            msg = cl.events[-1][1]
        if cl.transfer:
            msg = f"Downloading {cl.transfer} (Esc cancels)"
        line1 = f"{status}  {args.host}:{args.port}   input: {typed}_" + (f"   | {msg}" if msg else "")
        win.blit(font.render(line1, True, (230, 230, 230)), (8, H + 4))
        hint = "<- back  -> links  v/^ page  F1 index F2 home F3 files  End copy  ^R refresh  ^L go to  ^S shot  ^K scanlines  ^H help  ^Q quit"
        win.blit(small.render(hint, True, (140, 140, 150)), (8, H + 26))

        if ui.entry:                                   # text prompt overlay
            e = ui.entry
            pygame.draw.rect(win, (20, 40, 90), (0, H, W, BAR))
            win.blit(small.render(e["label"] + "   (Enter = OK, Esc = cancel)", True, (190, 210, 255)), (8, H + 4))
            win.blit(font.render(e["text"] + "_", True, (255, 255, 255)), (8, H + 22))
        if show_help:
            pygame.draw.rect(win, (0, 0, 0), (40, 40, W - 80, 40 + 28 * len(HELP)))
            for i, t in enumerate(HELP):
                win.blit(font.render(t, True, (255, 255, 120)), (60, 56 + 28 * i))
        pygame.display.flip()
        clock.tick(30)
    cl.close()
    pygame.quit()
    return 0


class HeadlessUI:
    def __init__(self, downloads):
        self.downloads = downloads

    def ask_save(self, desc, size):
        return os.path.join(self.downloads, safe_name(desc))

    def cancel_requested(self):
        return False


def run_headless(args):
    """Connect, optionally go to a page, wait for the line to go quiet, save a PNG."""
    screen = Screen()
    screen.scanlines = args.scanlines
    cl = Client(args.host, args.port, screen, HeadlessUI(args.downloads), args.baud)
    cl.connect()

    def settle():
        time.sleep(0.3)
        while cl.alive and time.monotonic() - cl.last_rx < args.settle:
            time.sleep(0.05)

    settle()                                   # welcome page
    cl.send("\r")                              # any first input -> index.bbc
    settle()
    if args.go:
        cl.send(args.go + "\r")
        settle()
    for k in args.keys or []:
        cl.send(k.encode().decode("unicode_escape").encode("latin-1"))
        settle()
    screen.render().save(args.shot)
    for t, m in cl.events:
        print("note:", m)
    print("saved", args.shot)
    cl.close()
    return 0


def main():
    ap = argparse.ArgumentParser(description="BeebWeb desktop client")
    ap.add_argument("host", nargs="?", default="beebs.bbs.io")
    ap.add_argument("port", nargs="?", type=int, default=8065)
    ap.add_argument("--scale", type=float, default=1, help="window scale (default 1.5)")
    ap.add_argument("--scanlines", type=float, default=0.4,
                    help="CRT scanline darkness 0-0.9 (default 0.4, 0 = off; Ctrl+K toggles)")
    ap.add_argument("--baud", type=int, default=0,
                    help="throttle to a serial speed, e.g. 19200, to feel real download times")
    ap.add_argument("--downloads", default="downloads", help="default folder for saved files")
    ap.add_argument("--shot", help="headless: save a PNG of the page and exit")
    ap.add_argument("--go", help="headless: page to load first, e.g. news.bbc")
    ap.add_argument("--keys", nargs="*", help="headless: extra keys to send, e.g. '\\x8e' '3\\r'")
    ap.add_argument("--settle", type=float, default=1.0, help="headless: seconds of quiet = page finished")
    args = ap.parse_args()
    sys.exit(run_headless(args) if args.shot else run_gui(args))


if __name__ == "__main__":
    main()