"""Stamp digests over core NATS. Subscribe to verae.ts.completed before requesting.

No status poll and no retry sleep. A known digest returns its certificate on the
request reply. A new digest returns a jobId, then verae.ts.completed carries the
certificate.

Usage:
  python scripts/nats_stamp.py <64-hex-or-text> [...]
  VERAE_NATS_URL=nats://127.0.0.1:4222 python scripts/nats_stamp.py "Hello World!"
"""

from __future__ import annotations

import hashlib
import json
import os
import socket
import sys
import time

SUBJECT_REQUEST = "verae.ts.request"
SUBJECT_COMPLETED = "verae.ts.completed"


def digest_of(text: str) -> str:
    raw = text.strip().lower()
    if raw.startswith("sha256:"):
        raw = raw[7:]
    if len(raw) == 64 and all(c in "0123456789abcdef" for c in raw):
        return raw
    return hashlib.sha256(text.encode()).hexdigest()


class Conn:
    def __init__(self, url: str) -> None:
        host_port = url.split("://", 1)[-1]
        host, port_s = host_port.rsplit(":", 1)
        self.sock = socket.create_connection((host, int(port_s)), timeout=5)
        self.buf = b""
        info = self._read_line()
        if not info.startswith("INFO "):
            raise RuntimeError(info)
        self.queued: dict[str, list[dict]] = {}
        self._send(
            'CONNECT {"verbose":false,"pedantic":false,"lang":"verae-stamp","version":"1","protocol":1}\r\n'
            "PING\r\n"
        )
        self._expect_pong()

    def _send(self, text: str) -> None:
        self.sock.sendall(text.encode())

    def _read_line(self) -> str:
        while b"\r\n" not in self.buf:
            chunk = self.sock.recv(65536)
            if not chunk:
                raise RuntimeError("nats closed")
            self.buf += chunk
        line, self.buf = self.buf.split(b"\r\n", 1)
        return line.decode()

    def _expect_pong(self) -> None:
        while True:
            line = self._read_line()
            if line == "PONG" or line.startswith("+OK"):
                if line == "PONG":
                    return
                continue
            if line.startswith("-ERR"):
                raise RuntimeError(line)

    def subscribe(self, subject: str, sid: str) -> None:
        self._send(f"SUB {subject} {sid}\r\n")

    def request(self, subject: str, payload: dict, inbox: str, sid: str) -> dict:
        raw = json.dumps(payload).encode()
        self.subscribe(inbox, sid)
        self._send(f"PUB {subject} {inbox} {len(raw)}\r\n")
        self._send(raw.decode() + "\r\n")
        return self._wait_sid(sid, 30.0)

    def _read_msg(self) -> tuple[str, dict]:
        while True:
            line = self._read_line()
            if line.startswith("PING"):
                self._send("PONG\r\n")
                continue
            if line.startswith("PONG") or line.startswith("+OK"):
                continue
            if not line.startswith("MSG "):
                continue
            parts = line.split()
            size = int(parts[-1])
            msg_sid = parts[2]
            while len(self.buf) < size + 2:
                chunk = self.sock.recv(65536)
                if not chunk:
                    raise RuntimeError("nats closed")
                self.buf += chunk
            body = self.buf[:size]
            self.buf = self.buf[size + 2 :]
            return msg_sid, json.loads(body.decode())

    def _wait_sid(self, sid: str, timeout: float) -> dict:
        queued = self.queued.get(sid) or []
        if queued:
            return queued.pop(0)
        deadline = time.time() + timeout
        while time.time() < deadline:
            self.sock.settimeout(max(0.1, deadline - time.time()))
            msg_sid, body = self._read_msg()
            if msg_sid == sid:
                return body
            self.queued.setdefault(msg_sid, []).append(body)
        raise TimeoutError(sid)

    def next_completed(self, timeout: float) -> dict:
        return self._wait_sid("done", timeout)


def main(argv: list[str]) -> int:
    if not argv:
        print("usage: nats_stamp.py <text-or-digest> [...]", file=sys.stderr)
        return 2
    url = os.environ.get("VERAE_NATS_URL", "nats://127.0.0.1:4222")
    digests = [digest_of(item) for item in argv]
    conn = Conn(url)
    conn.subscribe(SUBJECT_COMPLETED, "done")
    pending: dict[str, str] = {}
    receipts: dict[str, str] = {}
    t0 = time.perf_counter()
    for i, digest in enumerate(digests):
        started = time.perf_counter()
        body = conn.request(SUBJECT_REQUEST, {"hash": digest}, f"_INBOX.{i}", f"inbox{i}")
        if body.get("error"):
            print(json.dumps({"hash": digest, "error": body["error"]}))
            return 1
        elapsed = (time.perf_counter() - started) * 1000
        if body.get("receipt_id"):
            receipts[digest] = body["receipt_id"]
            print(json.dumps({"hash": digest, "receipt_id": body["receipt_id"], "existing": body.get("existing"), "ms": round(elapsed, 3)}))
        else:
            pending[body["jobId"]] = digest
            print(json.dumps({"hash": digest, "jobId": body.get("jobId"), "existing": False, "ms": round(elapsed, 3)}))
    while pending:
        msg = conn.next_completed(30)
        job_id = msg.get("jobId")
        digest = pending.pop(job_id, None)
        if digest is None:
            continue
        receipts[digest] = msg.get("receipt_id") or ""
        print(json.dumps({"hash": digest, "jobId": job_id, "receipt_id": receipts[digest], "via": "verae.ts.completed"}))
    print(json.dumps({"n": len(digests), "receipts": len(receipts), "wall_ms": round((time.perf_counter() - t0) * 1000, 3)}))
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))