/**
 * JSON-RPC connection over an injected write function.
 *
 * Provides sendRequest (correlated by id), sendNotification, onRequest,
 * 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 RequestHandler = (params: unknown) => unknown | Promise<unknown>;
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 JsonRpcConnection {
	private nextId = 1;
	private pending = new Map<number | string, PendingRequest>();
	private requestHandlers = new Map<string, RequestHandler>();
	private notificationHandlers = new Map<string, NotificationHandler>();
	private write: WriteFn;

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

	sendRequest(method: string, params?: unknown): Promise<unknown> {
		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);
		});
	}

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

	onRequest(method: string, handler: RequestHandler): void {
		this.requestHandlers.set(method, handler);
	}

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

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

		if (id !== undefined && method !== undefined) {
			await this.handleIncomingRequest(id, method, msg.params);
		} else if (method !== 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 async handleIncomingRequest(
		id: number | string,
		method: string,
		params: unknown,
	): Promise<void> {
		const handler = this.requestHandlers.get(method);
		if (!handler) {
			this.sendMessage({
				jsonrpc: "2.0",
				id,
				error: { code: -32601, message: `Method not found: ${method}` },
			});
			return;
		}
		try {
			const result = await handler(params);
			this.sendMessage({ jsonrpc: "2.0", id, result });
		} catch (err: unknown) {
			const message = err instanceof Error ? err.message : String(err);
			this.sendMessage({
				jsonrpc: "2.0",
				id,
				error: { code: -32603, message },
			});
		}
	}

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

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