// LeavePulse SDK — bidirectional realtime client.
//
// THE single WS contract between a client and the system. One multiplexed
// socket carries every topic subscription AND bidirectional exchanges:
//   - fan-out:        server → client `event` frames for subscribed topics
//   - request/reply:  client → server `request`, server → client `reply`
//   - session input:  client → server `input` (for future console/exec)
// The wire-format is protobuf (ws.proto: ClientFrame / ServerFrame), so the
// contract evolves by field tags without rewrites and is shared by every client
// (browser, node) and the gateway. This client owns connect/auth/reconnect; Vue
// or other UI layers wrap it, they never touch the socket.

import { create, fromBinary, toBinary } from "@bufbuild/protobuf";
import {
	ClientFrameSchema,
	type ServerFrame,
	ServerFrameSchema,
} from "./realtime_pb/leavepulse/realtime/v1/ws_pb.js";

/** Binary WebSocket shape both the browser and the `ws` package satisfy. */
export interface BinaryWebSocketLike {
	send(data: Uint8Array): void;
	close(): void;
	binaryType: string;
	addEventListener(type: "open" | "close" | "error", cb: () => void): void;
	addEventListener(type: "message", cb: (ev: { data: unknown }) => void): void;
}

export type RealtimeEventKind = "initial" | "update" | "error";

export interface RealtimeEnvelope<T = unknown> {
	topic: string;
	data: T;
}

export type RealtimeHandler<T = unknown> = (env: RealtimeEnvelope<T>) => void;

export interface RealtimeClientOptions {
	/** Build a socket for the given URL (browser `WebSocket`, or `ws`). */
	socketFactory: (url: string) => BinaryWebSocketLike;
	/** Base WS URL, e.g. `wss://rt.leavepulse.com`. */
	url: string;
	/** Async provider for the realtime auth (ws-)token; re-invoked on reconnect. */
	getToken?: () => Promise<string | null> | string | null;
	/** Reconnect backoff bounds (ms). */
	reconnectMinMs?: number;
	reconnectMaxMs?: number;
}

interface Subscription {
	topic: string;
	handlers: { kind: RealtimeEventKind; fn: RealtimeHandler }[];
	bootstrapped: boolean;
}

interface PendingRequest {
	resolve: (payload: Uint8Array) => void;
	reject: (err: Error) => void;
}

const DEFAULT_MIN_MS = 500;
const DEFAULT_MAX_MS = 30_000;
/**
 * How long before a ws-token expires to fetch and send the next one.
 *
 * The gateway accepts an `auth` frame at any point in a session and rebinds the
 * socket, but nothing was ever sending a second one: the token was presented
 * once at connect and then left to rot. A ws-token lives two minutes, so every
 * live subscription was torn down and rebuilt on that cycle — visible on the
 * server as a matched pair of "connection closed" every 120s, and to an
 * operator as live metrics that stall, because the agent's fast-sample tier is
 * driven by subscriber count and dropped to idle in each gap.
 *
 * The margin has to cover fetching a fresh token (a round trip to auth-service)
 * plus clock skew between the browser and the issuer.
 */
const AUTH_REFRESH_MARGIN_MS = 20_000;
/** Never re-arm faster than this, however short-lived the token claims to be. */
const AUTH_REFRESH_MIN_MS = 5_000;

/**
 * Milliseconds until a JWT's `exp`, or null if it carries no usable one.
 *
 * Deliberately does not verify the signature — the client cannot, and does not
 * need to: this only decides WHEN to ask for the next token. A malformed or
 * unsigned token yields null and refresh stays off, leaving the previous
 * reconnect behaviour rather than a tight loop.
 */
function millisUntilExpiry(token: string): number | null {
	// Destructured rather than indexed: consumers compile this file under their
	// own tsconfig, and one with noUncheckedIndexedAccess types parts[1] as
	// possibly undefined however the length was checked.
	const [header, payload, signature] = token.split(".");
	if (!header || !payload || !signature) return null;
	try {
		const base64 = payload.replace(/-/g, "+").replace(/_/g, "/");
		const padded = base64.padEnd(
			base64.length + ((4 - (base64.length % 4)) % 4),
			"=",
		);
		const claims = JSON.parse(atob(padded)) as { exp?: unknown };
		if (typeof claims.exp !== "number") return null;
		return claims.exp * 1000 - Date.now();
	} catch {
		return null;
	}
}

/** Decode a JSON payload carried in an Event/Reply `data` byte field. */
function decodeJson(bytes: Uint8Array): unknown {
	if (bytes.length === 0) return null;
	try {
		return JSON.parse(new TextDecoder().decode(bytes));
	} catch {
		return null;
	}
}

function encodeJson(value: unknown): Uint8Array {
	return new TextEncoder().encode(JSON.stringify(value ?? null));
}

/**
 * Owns the single multiplexed protobuf WebSocket. Reconnects with backoff,
 * re-authenticates and re-subscribes on reconnect, and routes events / replies.
 */
export class RealtimeClient {
	private socket: BinaryWebSocketLike | null = null;
	private connecting: Promise<void> | null = null;
	private authenticated = false;
	private closedByUser = false;
	private reconnectAttempts = 0;
	private requestSeq = 0;
	private readonly subscriptions = new Map<string, Subscription>();
	private readonly pending = new Map<string, PendingRequest>();
	// Resolved when the server's `welcome` frame lands — i.e. auth has been
	// processed. Subscribes wait on this so a private topic is never requested
	// before the gateway has bound the socket's identity (otherwise it replies
	// "Authentication required for private topic" and we needlessly reconnect).
	private welcomed: Promise<void> | null = null;
	private resolveWelcome: (() => void) | null = null;
	// Pending re-auth for the CURRENT socket. Cleared on close so a timer from a
	// dead socket cannot fire an auth frame into its replacement.
	private authRefresh: ReturnType<typeof setTimeout> | null = null;

	constructor(private readonly opts: RealtimeClientOptions) {}

	get isAuthenticated(): boolean {
		return this.authenticated;
	}

	/** Subscribe to a topic; returns an unsubscribe function. */
	async subscribe(
		topic: string,
		kind: RealtimeEventKind,
		handler: RealtimeHandler,
	): Promise<() => void> {
		await this.ensureConnected();
		let sub = this.subscriptions.get(topic);
		if (!sub) {
			sub = { topic, handlers: [], bootstrapped: false };
			this.subscriptions.set(topic, sub);
			this.sendSubscribe(topic);
		}
		sub.handlers.push({ kind, fn: handler });
		return () => this.removeHandler(topic, handler);
	}

	/** Send a request and await its single reply (the JSON-decoded payload). */
	async request<T = unknown>(method: string, payload?: unknown): Promise<T> {
		await this.ensureConnected();
		const id = `r${++this.requestSeq}`;
		const frame = create(ClientFrameSchema, {
			body: {
				case: "request",
				value: { id, method, payload: encodeJson(payload) },
			},
		});
		const result = new Promise<Uint8Array>((resolve, reject) => {
			this.pending.set(id, { resolve, reject });
		});
		this.sendFrame(frame);
		const bytes = await result;
		return decodeJson(bytes) as T;
	}

	/** Send a chunk of session input (e.g. console stdin) for a session. */
	async sendInput(sessionId: string, data: Uint8Array): Promise<void> {
		await this.ensureConnected();
		this.sendFrame(
			create(ClientFrameSchema, {
				body: { case: "input", value: { sessionId, data } },
			}),
		);
	}

	close(): void {
		this.closedByUser = true;
		// Explicit close bypasses onClose (we drop the socket ourselves), so the
		// refresh timer has to be cancelled here too or it outlives the client.
		this.clearAuthRefresh();
		this.subscriptions.clear();
		for (const p of this.pending.values())
			p.reject(new Error("realtime client closed"));
		this.pending.clear();
		this.socket?.close();
		this.socket = null;
		this.connecting = null;
		this.authenticated = false;
		this.resolveWelcome?.();
		this.resolveWelcome = null;
		this.welcomed = null;
	}

	// ── internals ────────────────────────────────────────────────────────────

	private removeHandler(topic: string, handler: RealtimeHandler): void {
		const sub = this.subscriptions.get(topic);
		if (!sub) return;
		sub.handlers = sub.handlers.filter((h) => h.fn !== handler);
		if (sub.handlers.length === 0) {
			this.subscriptions.delete(topic);
			this.sendFrame(
				create(ClientFrameSchema, {
					body: { case: "unsubscribe", value: { topic } },
				}),
			);
		}
	}

	private sendSubscribe(topic: string): void {
		this.sendFrame(
			create(ClientFrameSchema, {
				body: { case: "subscribe", value: { topic } },
			}),
		);
	}

	private async ensureConnected(): Promise<void> {
		if (this.socket && this.authenticated) return;
		if (this.connecting) return this.connecting;
		this.connecting = this.connect();
		return this.connecting;
	}

	private async connect(): Promise<void> {
		this.closedByUser = false;
		const token = this.opts.getToken ? await this.opts.getToken() : null;
		const socket = this.opts.socketFactory(this.opts.url);
		socket.binaryType = "arraybuffer";
		this.socket = socket;

		// Arm the welcome gate before sending auth, so the frame can't race us.
		this.welcomed = new Promise<void>((resolve) => {
			this.resolveWelcome = resolve;
		});

		await new Promise<void>((resolve, reject) => {
			socket.addEventListener("open", () => resolve());
			socket.addEventListener("error", () =>
				reject(new Error("realtime socket error")),
			);
		});
		socket.addEventListener("message", (ev) => this.onMessage(ev.data));
		socket.addEventListener("close", () => this.onClose());

		// Authenticate first so private subscriptions are accepted.
		if (token) {
			this.sendFrame(
				create(ClientFrameSchema, {
					body: { case: "auth", value: { token } },
				}),
			);
			this.scheduleAuthRefresh(token);
		}
		// Wait for the gateway's `welcome` (auth applied) before (re)subscribing,
		// so private topics aren't requested on an un-bound socket. The gate is
		// resolved in onMessage; if the socket drops first, onClose rejects the
		// retry path, so a hung await can't wedge us.
		await this.welcomed;
		// Re-subscribe everything (covers reconnect).
		for (const topic of this.subscriptions.keys()) this.sendSubscribe(topic);
		this.reconnectAttempts = 0;
	}

	/**
	 * Send a fresh token over the OPEN socket shortly before this one expires.
	 *
	 * Re-authenticating in place is what keeps a long-lived subscription alive:
	 * the alternative is letting the token lapse and being reconnected, which
	 * drops every subscription for the duration and, for subscriber-driven
	 * sources, stops the upstream feed entirely until it is re-established.
	 */
	private scheduleAuthRefresh(token: string): void {
		this.clearAuthRefresh();
		if (!this.opts.getToken) return;
		const remaining = millisUntilExpiry(token);
		if (remaining === null) return;
		const delay = Math.max(
			AUTH_REFRESH_MIN_MS,
			remaining - AUTH_REFRESH_MARGIN_MS,
		);
		this.authRefresh = setTimeout(() => {
			void this.refreshAuth();
		}, delay);
	}

	private async refreshAuth(): Promise<void> {
		const socket = this.socket;
		if (!socket || !this.opts.getToken) return;
		let next: string | null = null;
		try {
			next = await this.opts.getToken();
		} catch {
			next = null;
		}
		// Nothing to send, or the socket died while we were fetching: leave it to
		// the reconnect path rather than authenticating a socket that is gone.
		if (!next || this.socket !== socket) return;
		this.sendFrame(
			create(ClientFrameSchema, {
				body: { case: "auth", value: { token: next } },
			}),
		);
		this.scheduleAuthRefresh(next);
	}

	private clearAuthRefresh(): void {
		if (this.authRefresh === null) return;
		clearTimeout(this.authRefresh);
		this.authRefresh = null;
	}

	private onClose(): void {
		this.clearAuthRefresh();
		this.socket = null;
		this.connecting = null;
		this.authenticated = false;
		// Unblock a connect() still awaiting `welcome` so its promise can settle
		// (it'll re-arm on the next attempt); reconnect drives the retry.
		this.resolveWelcome?.();
		this.resolveWelcome = null;
		if (this.closedByUser) return;
		this.scheduleReconnect();
	}

	private scheduleReconnect(): void {
		const min = this.opts.reconnectMinMs ?? DEFAULT_MIN_MS;
		const max = this.opts.reconnectMaxMs ?? DEFAULT_MAX_MS;
		const delay = Math.min(max, min * 2 ** this.reconnectAttempts++);
		setTimeout(() => {
			if (this.closedByUser) return;
			void this.ensureConnected().catch(() => this.scheduleReconnect());
		}, delay);
	}

	private onMessage(raw: unknown): void {
		const bytes = this.toBytes(raw);
		if (!bytes) return;
		let frame: ServerFrame;
		try {
			frame = fromBinary(ServerFrameSchema, bytes);
		} catch {
			return;
		}
		switch (frame.body.case) {
			case "ping":
				this.sendFrame(
					create(ClientFrameSchema, { body: { case: "ping", value: {} } }),
				);
				return;
			case "welcome":
				this.authenticated = frame.body.value.authenticated;
				// Release any subscribes waiting on the auth handshake.
				this.resolveWelcome?.();
				this.resolveWelcome = null;
				return;
			case "event": {
				const ev = frame.body.value;
				this.dispatchEvent(ev.topic, decodeJson(ev.data));
				return;
			}
			case "reply": {
				const reply = frame.body.value;
				const pending = this.pending.get(reply.id);
				if (!pending) return;
				this.pending.delete(reply.id);
				if (reply.error)
					pending.reject(
						new Error(
							reply.error.message || reply.error.code || "request failed",
						),
					);
				else pending.resolve(reply.payload);
				return;
			}
			case "error": {
				const err = frame.body.value;
				if (err.id && this.pending.has(err.id)) {
					const p = this.pending.get(err.id);
					this.pending.delete(err.id);
					p?.reject(new Error(err.message || err.code || "realtime error"));
					return;
				}
				if (err.topic) this.dispatchError(err.topic, err);
				return;
			}
			default:
				return;
		}
	}

	private toBytes(raw: unknown): Uint8Array | null {
		if (raw instanceof Uint8Array) return raw;
		if (raw instanceof ArrayBuffer) return new Uint8Array(raw);
		return null;
	}

	private dispatchEvent(topic: string, data: unknown): void {
		const sub = this.subscriptions.get(topic);
		if (!sub) return;
		const kind: RealtimeEventKind = sub.bootstrapped ? "update" : "initial";
		sub.bootstrapped = true;
		for (const h of sub.handlers) if (h.kind === kind) h.fn({ topic, data });
	}

	private dispatchError(topic: string, error: unknown): void {
		const sub = this.subscriptions.get(topic);
		if (!sub) return;
		for (const h of sub.handlers)
			if (h.kind === "error") h.fn({ topic, data: error });
	}

	private sendFrame(
		frame: ReturnType<typeof create<typeof ClientFrameSchema>>,
	): void {
		if (!this.socket) return;
		this.socket.send(toBinary(ClientFrameSchema, frame));
	}
}