#!/usr/bin/env python3
"""Raw protocol v3 client (no libpq).  It records the exact bytes of every
message the server sends, i.e. what is on the socket, with no dissector in
between.

  rfq_wire_client.py capture <sockdir> <port> [options]
      -> JSON with every server message for three phases: startup, a simple
         Query 'SELECT 1', and an extended Parse/Bind/Execute/Sync 'SELECT 1'.
  rfq_wire_client.py script <sockdir> <port> [options] -- 'sql1' 'sql2' ...
      -> for each statement (sent as a simple Query), the 'Z' message that
         closes it, raw and decoded, plus the first row and any error.

  [options] is the value of the startup 'options' parameter,
  e.g. '-c ready_for_query_message=rich'.

Connects over a Unix socket as user postgres and expects trust authentication.
"""
import json
import socket
import struct
import sys


class Conn:
    def __init__(self, sockdir, port, options=None, user="postgres", db="postgres"):
        self.s = socket.socket(socket.AF_UNIX)
        self.s.connect(f"{sockdir}/.s.PGSQL.{port}")
        self.buf = b""
        params = [b"user", user.encode(), b"database", db.encode()]
        if options:
            params += [b"options", options.encode()]
        body = struct.pack("!I", 196608) + b"".join(p + b"\0" for p in params) + b"\0"
        self.s.sendall(struct.pack("!I", len(body) + 4) + body)

    def _read(self, n):
        while len(self.buf) < n:
            chunk = self.s.recv(65536)
            if not chunk:
                raise EOFError("server closed the connection")
            self.buf += chunk
        out, self.buf = self.buf[:n], self.buf[n:]
        return out

    def msg(self):
        head = self._read(5)
        typ, length = head[:1], struct.unpack("!I", head[1:])[0]
        return typ, head + self._read(length - 4)

    def until_z(self):
        msgs = []
        while True:
            t, raw = self.msg()
            msgs.append((t, raw))
            if t == b"R" and struct.unpack("!I", raw[5:9])[0] != 0:
                raise RuntimeError("server asks for authentication; use trust")
            if t == b"Z":
                return msgs

    def send(self, typ, body):
        self.s.sendall(typ + struct.pack("!I", len(body) + 4) + body)

    def query(self, sql):
        self.send(b"Q", sql.encode() + b"\0")
        return self.until_z()

    def extended(self, sql):
        self.send(b"P", b"\0" + sql.encode() + b"\0" + struct.pack("!H", 0))
        self.send(b"B", b"\0\0" + struct.pack("!HHH", 0, 0, 0))
        self.send(b"E", b"\0" + struct.pack("!I", 0))
        self.send(b"S", b"")
        return self.until_z()


def decode_z(raw):
    """Raw 'Z' -> {len, status, kv: [(key, value)], unparsed}."""
    length = struct.unpack("!I", raw[1:5])[0]
    body = raw[5:]
    out = {"len": length, "status": body[:1].decode(), "kv": []}
    off = 1
    while off < len(body):
        klen = body[off]; off += 1
        k = body[off:off + klen]; off += klen
        if off >= len(body):
            out["unparsed"] = body[off - klen - 1:].hex(); break
        vlen = body[off]; off += 1
        v = body[off:off + vlen]; off += vlen
        out["kv"].append([k.decode(), v.decode()])
    return out


def describe(t, raw):
    d = {"type": t.decode(), "len": len(raw), "hex": raw.hex()}
    if t == b"S":
        k, v, _ = raw[5:].split(b"\0", 2)
        d["param"] = [k.decode(), v.decode()]
    if t == b"K":
        d["note"] = "BackendKeyData: differs on every connection, ignore when comparing"
    if t == b"Z":
        d["z"] = decode_z(raw)
    return d


def main():
    mode, sockdir, port = sys.argv[1], sys.argv[2], sys.argv[3]
    rest = sys.argv[4:]
    options = None
    if rest and rest[0] != "--":
        options = rest.pop(0)
    c = Conn(sockdir, port, options)
    if mode == "capture":
        phases = {"startup": c.until_z(), "simple_select_1": c.query("SELECT 1"),
                  "extended_select_1": c.extended("SELECT 1")}
        print(json.dumps({p: [describe(t, r) for t, r in m] for p, m in phases.items()}, indent=1))
    elif mode == "script":
        sqls = rest[1:] if rest and rest[0] == "--" else rest
        res = [{"sql": "(startup)", "z": decode_z(c.until_z()[-1][1])}]
        for sql in sqls:
            msgs = c.query(sql)
            err = [r[5:].split(b"\0M")[1].split(b"\0")[0].decode() for t, r in msgs if t == b"E"]
            rows = [r for t, r in msgs if t == b"D"]
            row = None
            if rows:
                r = rows[0]; n = struct.unpack("!H", r[5:7])[0]; off = 7; vals = []
                for _ in range(n):
                    vlen = struct.unpack("!i", r[off:off + 4])[0]; off += 4
                    vals.append(None if vlen < 0 else r[off:off + vlen].decode()); off += max(vlen, 0)
                row = vals
            res.append({"sql": sql, "row": row, "error": err or None,
                        "z_hex": msgs[-1][1].hex(), "z": decode_z(msgs[-1][1])})
        print(json.dumps(res, indent=1))
    else:
        sys.exit(__doc__)


if __name__ == "__main__":
    main()