#!/usr/bin/env python3
"""Stdlib-only Home Assistant WebSocket client (via the Supervisor core proxy).

Why this exists
---------------
Some HA operations have **no REST route** and must go through the WebSocket API
-- writing Lovelace config (`lovelace/config/save`) is the main one, and the
`.storage/` files must not be hand-edited (HA may overwrite them, and edits skip
validation). The obvious paths are dead ends on this host: the HA host and the
core container have no usable WS client (the Supervisor proxy rejects core's own
loop), and the Supervisor image has neither `websockets` nor `aiohttp`.
HA's WS protocol is plain text frames, so ~90 lines of stdlib is enough.

Run it from the HA host (see hosts/hass.windy.lan.md "改法（可复用）"):

    B64=$(base64 < scripts/ha-ws-client.py | tr -d '\n')
    echo '{"action":"get","url_path":"dashboard-quick"}' | \
      docker run --rm -i --network host -e SUPERVISOR_TOKEN --entrypoint python3 \
        r.hassbus.com/home-assistant/aarch64-hassio-supervisor:<ver> \
        -c "import base64,sys;exec(base64.b64decode('$B64').decode())" - -

Do NOT use `-v /tmp/...`: that path is resolved by the Docker daemon on the
HAOS host, not inside the SSH add-on, so the mount comes up empty.

Payload / output
----------------
    {"action":"get",  "url_path":"dashboard-quick"}
    {"action":"save", "url_path":"dashboard-quick", "config": {...}}
    {"action":"raw",  "command": {"type":"lovelace/resources/list"}}

    python3 ha-ws-client.py <payload.json|-> <out.json|->

"-" reads the payload from stdin / prints the result to stdout, wrapped in
<<<DSH_RESULT>>> / <<<DSH_END>>> markers so it survives the add-on login banner.
"""
import base64
import json
import os
import random
import socket
import struct
import sys

HOST = os.environ.get("WS_HOST", "172.30.32.2")
PORT = int(os.environ.get("WS_PORT", "80"))
PATH = os.environ.get("WS_PATH", "/core/websocket")


class WS:
    def __init__(self, host, port, path, extra_headers=None):
        self.s = socket.create_connection((host, port), timeout=30)
        key = base64.b64encode(bytes(random.getrandbits(8) for _ in range(16))).decode()
        req = (
            f"GET {path} HTTP/1.1\r\nHost: {host}\r\nUpgrade: websocket\r\n"
            f"Connection: Upgrade\r\nSec-WebSocket-Key: {key}\r\n"
            f"Sec-WebSocket-Version: 13\r\n"
        )
        for k, v in (extra_headers or {}).items():
            req += f"{k}: {v}\r\n"
        self.s.sendall((req + "\r\n").encode())
        buf = b""
        while b"\r\n\r\n" not in buf:
            d = self.s.recv(4096)
            if not d:
                raise RuntimeError("closed during handshake")
            buf += d
        head, _, rest = buf.partition(b"\r\n\r\n")
        status = head.split(b"\r\n")[0]
        if b"101" not in status:
            raise RuntimeError("handshake failed: " + status.decode(errors="replace"))
        self.buf = rest

    def _exact(self, n):
        while len(self.buf) < n:
            d = self.s.recv(65536)
            if not d:
                raise RuntimeError("socket closed")
            self.buf += d
        out, self.buf = self.buf[:n], self.buf[n:]
        return out

    def send(self, text):
        data = text.encode()
        mask = bytes(random.getrandbits(8) for _ in range(4))
        n = len(data)
        hdr = bytearray([0x81])
        if n < 126:
            hdr.append(0x80 | n)
        elif n < 65536:
            hdr.append(0x80 | 126)
            hdr += struct.pack(">H", n)
        else:
            hdr.append(0x80 | 127)
            hdr += struct.pack(">Q", n)
        hdr += mask
        self.s.sendall(bytes(hdr) + bytes(b ^ mask[i % 4] for i, b in enumerate(data)))

    def recv(self):
        while True:
            b0, b1 = self._exact(2)
            op = b0 & 0x0F
            ln = b1 & 0x7F
            if ln == 126:
                ln = struct.unpack(">H", self._exact(2))[0]
            elif ln == 127:
                ln = struct.unpack(">Q", self._exact(8))[0]
            payload = self._exact(ln) if ln else b""
            if op == 0x8:
                raise RuntimeError("server closed")
            if op in (0x9, 0xA):
                continue
            return payload.decode()


def main():
    token = os.environ["SUPERVISOR_TOKEN"]
    src = sys.argv[1]
    payload = json.load(sys.stdin if src == "-" else open(src, encoding="utf-8"))
    action = payload.pop("action")
    url_path = payload.pop("url_path", "dashboard-quick")
    out_path = sys.argv[2]

    ws = WS(HOST, PORT, PATH, {"Authorization": f"Bearer {token}"})
    while True:
        hello = json.loads(ws.recv())
        if hello.get("type") == "auth_required":
            break
        if hello.get("type") == "auth_ok":
            break
    ws.send(json.dumps({"type": "auth", "access_token": token}))
    while True:
        r = json.loads(ws.recv())
        if r.get("type") == "auth_ok":
            break
        if r.get("type") == "auth_invalid":
            print("AUTH_INVALID:", json.dumps(r, ensure_ascii=False)[:300])
            return 1

    cmd = {
        "id": 1,
        "type": "lovelace/config/save" if action == "save" else "lovelace/config",
        "url_path": url_path,
    }
    if action == "save":
        cmd["config"] = payload["config"]
    elif action == "raw":
        cmd = dict(payload["command"])
        cmd["id"] = 1
    ws.send(json.dumps(cmd))
    while True:
        r = json.loads(ws.recv())
        if r.get("id") != 1:
            continue
        if not r.get("success"):
            print("ERROR:", json.dumps(r, ensure_ascii=False)[:600])
            return 1
        result = r.get("result")
        blob = json.dumps(result if result is not None else {"ok": True},
                          ensure_ascii=False, indent=1)
        if out_path == "-":
            print("<<<DSH_RESULT>>>")
            print(blob)
            print("<<<DSH_END>>>")
        else:
            with open(out_path, "w", encoding="utf-8") as f:
                f.write(blob)
            print("OK ->", out_path)
        return 0


if __name__ == "__main__":
    sys.exit(main())