#!/usr/bin/env python3
"""Import the status page (19) or the Restart page (20; both numbers come from config.json) into Companion without clicking through the UI.

Does exactly what Import/Export -> Import -> "import into page N" does, over the
same internal API the Companion web UI uses (tRPC over a WebSocket at /trpc):
upload the file in chunks, check its SHA-1, then import it as a single page.

Safety: it looks at what is on the target page, not just its name. It refuses unless every
key on the page (besides Companion's own page up / number / down keys) carries one of the
markers make_page.py writes into its notes, i.e. the page holds nothing a person made by hand.
A key someone added is named by its position, so it can be moved first.

  python3 push_page.py                         # the status page, from config.json
  python3 push_page.py --which restart         # the Restart page
  python3 push_page.py --host 10.6.3.222 --page 19 --file status-page.companionconfig

Stdlib only (a minimal WebSocket client is included). Written against Companion 5.0.7.
"""
import argparse
import base64
import gzip
import hashlib
import json
import os
import socket
import struct
import sys
import time
import urllib.parse
import urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from make_page import MARKERS  # noqa: E402

CHUNK = 256 * 1024


class WebSocket:
    def __init__(self, host, port, path, timeout=15):
        self.sock = socket.create_connection((host, port), timeout=timeout)
        key = base64.b64encode(os.urandom(16)).decode()
        self.sock.sendall(("GET %s HTTP/1.1\r\nHost: %s:%d\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n"
                           "Sec-WebSocket-Key: %s\r\nSec-WebSocket-Version: 13\r\n\r\n"
                           % (path, host, port, key)).encode())
        buf = b""
        while b"\r\n\r\n" not in buf:
            chunk = self.sock.recv(4096)
            if not chunk:
                raise ConnectionError("connection closed during WebSocket handshake")
            buf += chunk
        head, self.buf = buf.split(b"\r\n\r\n", 1)
        if b" 101 " not in head.split(b"\r\n", 1)[0]:
            raise ConnectionError("WebSocket handshake refused: %s" % head.split(b"\r\n", 1)[0].decode())

    def _read(self, n):
        while len(self.buf) < n:
            chunk = self.sock.recv(65536)
            if not chunk:
                raise ConnectionError("connection closed")
            self.buf += chunk
        data, self.buf = self.buf[:n], self.buf[n:]
        return data

    def _send(self, opcode, payload):
        n = len(payload)
        header = bytes([0x80 | opcode])
        if n < 126:
            header += bytes([0x80 | n])
        elif n < 65536:
            header += bytes([0x80 | 126]) + struct.pack(">H", n)
        else:
            header += bytes([0x80 | 127]) + struct.pack(">Q", n)
        mask = os.urandom(4)
        masked = bytes(b ^ mask[i & 3] for i, b in enumerate(payload))
        self.sock.sendall(header + mask + masked)

    def send(self, text):
        self._send(0x1, text.encode())

    def recv(self):
        message = b""
        while True:
            b1, b2 = self._read(2)
            fin, opcode, n = b1 & 0x80, b1 & 0x0F, b2 & 0x7F
            if n == 126:
                n = struct.unpack(">H", self._read(2))[0]
            elif n == 127:
                n = struct.unpack(">Q", self._read(8))[0]
            payload = self._read(n)
            if opcode == 0x9:                       # ping
                self._send(0xA, payload)
                continue
            if opcode == 0x8:
                raise ConnectionError("server closed the WebSocket")
            message += payload
            if fin:
                return message.decode()

    def close(self):
        try:
            self._send(0x8, b"")
        finally:
            self.sock.close()


class Trpc:
    def __init__(self, ws):
        self.ws, self.next_id = ws, 1

    def mutation(self, path, data):
        rid = self.next_id
        self.next_id += 1
        self.ws.send(json.dumps({"id": rid, "jsonrpc": "2.0", "method": "mutation",
                                 "params": {"path": path, "input": data}}))
        while True:
            msg = json.loads(self.ws.recv())
            if msg.get("id") != rid:
                continue
            if "error" in msg:
                err = msg["error"]
                raise RuntimeError("%s failed: %s" % (path, err.get("message", err)))
            return msg.get("result", {}).get("data")


NAV_TYPES = ("pageup", "pagedown", "pagenum")
NAV_MARKER = "STATUS nav key"          # our coloured page number key (make_page.nav_key)


def fetch_page(host, port, page):
    """The page as Companion exports it ({"page": {"name", "controls": {row: {col: control}}}}),
    without any connection secrets."""
    url = "http://%s:%d/int/export/page/%d?includeSecrets=false" % (host, port, page)
    with urllib.request.urlopen(url, timeout=10) as r:
        raw = r.read()
    try:
        raw = gzip.decompress(raw)
    except OSError:
        pass
    return json.loads(raw)


def foreign_buttons(export):
    """[(row, col)] of every key on the page that neither Companion nor this tool made."""
    page = export.get("page") or {}
    found = []
    for r, row in sorted((page.get("controls") or {}).items(), key=lambda kv: int(kv[0])):
        for c, ctl in sorted((row or {}).items(), key=lambda kv: int(kv[0])):
            if not ctl or ctl.get("type") in NAV_TYPES:
                continue
            notes = (ctl.get("options") or {}).get("notes") or ""
            if any(m in notes for m in MARKERS):
                continue
            found.append((int(r), int(c)))
    return found


def real_buttons(export):
    """How many keys a page holds that this tool did not make."""
    return len(foreign_buttons(export))


def safe_target(export, expected_name=None):
    """(ok, reason). OK only when every key on the page is Companion's own page key or one of ours,
    whatever the page is called: nothing a person made by hand is ever overwritten."""
    name = ((export.get("page") or {}).get("name") or "")
    foreign = foreign_buttons(export)
    if not foreign:
        return True, "it holds only page keys and keys made by this tool"
    where = ", ".join("row %d col %d" % rc for rc in foreign[:6]) + (" ..." if len(foreign) > 6 else "")
    return False, "it is called %r and holds %d key%s not made by this tool (%s) - move or delete them first" % (
        name, len(foreign), "" if len(foreign) == 1 else "s", where)


def page_name(host, port, page):
    url = "http://%s:%d/api/variable/internal/page_number_%d_name/value" % (host, port, page)
    with urllib.request.urlopen(url, timeout=5) as r:
        return r.read().decode()


def push(host, port, page, file_path, expected_name):
    """Upload a built page file and import it into `page` (the caller has checked safe_target)."""
    data = open(file_path, "rb").read()
    built = json.loads(gzip.decompress(data))
    if (built.get("page") or {}).get("name") != expected_name:
        raise ValueError("%s builds page %r, not %r" % (os.path.basename(file_path), (built.get("page") or {}).get("name"), expected_name))
    ws = WebSocket(host, port, "/trpc")
    try:
        api = Trpc(ws)
        for attempt in range(40):
            try:
                session = api.mutation("importExport.prepareImport.start", {"name": os.path.basename(file_path), "size": len(data)})
                break
            except RuntimeError as e:
                if "already exists" not in str(e) or attempt == 39:
                    raise
                time.sleep(3)
        for offset in range(0, len(data), CHUNK):
            api.mutation("importExport.prepareImport.uploadChunk",
                         {"sessionId": session, "offset": offset, "data": base64.b64encode(data[offset:offset + CHUNK]).decode()})
        summary = api.mutation("importExport.prepareImport.complete",
                               {"sessionId": session, "expectedChecksum": hashlib.sha1(data).hexdigest(), "userData": None})
        if isinstance(summary, list) and summary and summary[0]:
            raise RuntimeError("Companion could not read the file: %s" % summary[0])
        api.mutation("importExport.importSinglePage", {"targetPage": page, "sourcePage": 1, "connectionIdRemapping": {}})
    finally:
        ws.close()


def main():
    cfg = json.load(open(os.path.join(HERE, "config.json")))
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--which", choices=("status", "restart"), default="status")
    ap.add_argument("--host", default=cfg.get("site", {}).get("ingest_hostname") or cfg["companion"].get("host", "127.0.0.1"))
    ap.add_argument("--port", type=int, default=cfg["companion"].get("port", 8000))
    ap.add_argument("--page", type=int)
    ap.add_argument("--file")
    ap.add_argument("--force", action="store_true", help="overwrite the page whatever is on it")
    args = ap.parse_args()
    if args.which == "restart":
        page, expected = cfg["restart_page"]["page"], cfg["restart_page"]["page_name"]
        default_file = "restart-page.companionconfig"
    else:
        page, expected, default_file = cfg["page"], cfg["page_name"], "status-page.companionconfig"
    args.page = args.page or page
    args.file = args.file or os.path.join(HERE, default_file)

    data = open(args.file, "rb").read()
    built = json.loads(gzip.decompress(data))
    if (built.get("page") or {}).get("name") != expected:
        sys.exit("%s builds page %r, not %r; refusing" % (os.path.basename(args.file),
                                                          (built.get("page") or {}).get("name"), expected))
    ok, why = safe_target(fetch_page(args.host, args.port, args.page), expected)
    if not ok and not args.force:
        sys.exit("page %d: %s - refusing (use --force only if you mean to replace it)" % (args.page, why))
    plain = page_name(args.host, args.port, args.page).replace("\\n", " ").strip()
    ws = WebSocket(args.host, args.port, "/trpc")
    try:
        api = Trpc(ws)
        # Only one upload at a time; an abandoned one expires after a short idle time.
        for attempt in range(40):
            try:
                session = api.mutation("importExport.prepareImport.start",
                                       {"name": os.path.basename(args.file), "size": len(data)})
                break
            except RuntimeError as e:
                if "already exists" not in str(e) or attempt == 39:
                    raise
                if attempt == 0:
                    print("another upload is still open in Companion; waiting for it to expire...")
                time.sleep(3)
        for offset in range(0, len(data), CHUNK):
            api.mutation("importExport.prepareImport.uploadChunk",
                         {"sessionId": session, "offset": offset,
                          "data": base64.b64encode(data[offset:offset + CHUNK]).decode()})
        summary = api.mutation("importExport.prepareImport.complete",
                               {"sessionId": session, "expectedChecksum": hashlib.sha1(data).hexdigest(),
                                "userData": None})
        if isinstance(summary, list) and summary and summary[0]:
            sys.exit("Companion could not read the file: %s" % summary[0])
        api.mutation("importExport.importSinglePage",
                     {"targetPage": args.page, "sourcePage": 1, "connectionIdRemapping": {}})
    finally:
        ws.close()
    print("imported %s into page %d (was %r, now %r)" % (os.path.basename(args.file), args.page, plain,
                                                         page_name(args.host, args.port, args.page)))


if __name__ == "__main__":
    main()