/**
 * JSON-RPC 2.0 client over an injected write function.
 *
 * Provides request (correlated by id), notify, and onNotification.
 * The caller feeds decoded JSON messages via `handleMessage`.
 */

import { encode } from "./framing.js";

export type WriteFn = (bytes: Uint8Array) => void;

export interface PendingRequest {
	readonly resolve: (value: unknown) => void;
	readonly reject: (reason: unknown) => void;
}

export type NotificationHandler = (params: unknown) => void;

export interface JsonRpcMessage {
	readonly jsonrpc: "2.0";
	readonly id?: number | string | undefined;
	readonly method?: string | undefined;
	readonly params?: unknown;
	readonly result?: unknown;
	readonly error?:
		| { readonly code: number; readonly message: string; readonly data?: unknown }
		| undefined;
}

export class JsonRpcClient {
	private nextId = 1;
	private pending = new Map<number | string, PendingRequest>();
	private notificationHandlers = new Map<string, NotificationHandler>();
	private write: WriteFn;
	private closed = false;

	constructor(write: WriteFn) {
		this.write = write;
	}

	request(method: string, params?: unknown): Promise<unknown> {
		if (this.closed) {
			return Promise.reject(new Error("Connection closed"));
		}
		const id = this.nextId++;
		const msg: JsonRpcMessage = { jsonrpc: "2.0", id, method, params };
		return new Promise((resolve, reject) => {
			this.pending.set(id, { resolve, reject });
			this.sendMessage(msg);
		});
	}

	notify(method: string, params?: unknown): void {
		if (this.closed) return;
		const msg: JsonRpcMessage = { jsonrpc: "2.0", method, params };
		this.sendMessage(msg);
	}

	onNotification(method: string, handler: NotificationHandler): void {
		this.notificationHandlers.set(method, handler);
	}

	handleMessage(json: string): void {
		const msg = JSON.parse(json) as JsonRpcMessage;
		const { id, method } = msg;

		if (method !== undefined && id === undefined) {
			this.handleIncomingNotification(method, msg.params);
		} else if (id !== undefined) {
			this.handleResponse(id, msg);
		}
	}

	private sendMessage(msg: JsonRpcMessage): void {
		this.write(encode(JSON.stringify(msg)));
	}

	private handleResponse(id: number | string, msg: JsonRpcMessage): void {
		const entry = this.pending.get(id);
		if (!entry) return;
		this.pending.delete(id);
		if (msg.error) {
			entry.reject(new Error(msg.error.message));
		} else {
			entry.resolve(msg.result);
		}
	}

	private handleIncomingNotification(method: string, params: unknown): void {
		const handler = this.notificationHandlers.get(method);
		if (handler) {
			handler(params);
		}
	}

	close(): void {
		this.closed = true;
		for (const entry of this.pending.values()) {
			entry.reject(new Error("Connection closed"));
		}
		this.pending.clear();
	}
}