import { createConnection, type Socket } from 'node:net';
import { ReadBuffer, serializeMessage } from '@modelcontextprotocol/sdk/shared/stdio.js';
import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
import type { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';

/** MCP JSONL transport for an explicitly configured Unix-domain socket. */
export class UnixSocketClientTransport implements Transport {
  private socket?: Socket;
  private readonly readBuffer = new ReadBuffer();

  onclose?: () => void;
  onerror?: (error: Error) => void;
  onmessage?: (message: JSONRPCMessage) => void;

  constructor(private readonly socketPath: string) {}

  async start(): Promise<void> {
    if (this.socket) {
      throw new Error('UnixSocketClientTransport already started');
    }

    await new Promise<void>((resolve, reject) => {
      const socket = createConnection(this.socketPath);

      this.socket = socket;
      let connected = false;

      socket.once('connect', () => {
        connected = true;
        resolve();
      });
      socket.on('data', (chunk) => {
        try {
          this.readBuffer.append(typeof chunk === 'string' ? Buffer.from(chunk) : chunk);

          while (true) {
            const message = this.readBuffer.readMessage();

            if (message === null) {
              break;
            }

            this.onmessage?.(message);
          }
        } catch (error) {
          const cause = error instanceof Error ? error : new Error(String(error));

          this.onerror?.(cause);
          void this.close();
        }
      });
      socket.on('error', (error) => {
        if (!connected) {
          reject(error);
        }

        this.onerror?.(error);
      });
      socket.on('close', () => {
        if (this.socket === socket) {
          this.socket = undefined;
        }

        this.readBuffer.clear();
        this.onclose?.();
      });
    });
  }

  async close(): Promise<void> {
    const socket = this.socket;

    this.socket = undefined;
    this.readBuffer.clear();

    if (!socket || socket.destroyed) {
      return;
    }

    await new Promise<void>((resolve) => {
      const timeout = setTimeout(() => socket.destroy(), 2_000);

      timeout.unref();
      socket.once('close', () => {
        clearTimeout(timeout);
        resolve();
      });
      socket.end();
    });
  }

  async send(message: JSONRPCMessage): Promise<void> {
    const socket = this.socket;

    if (!socket || socket.destroyed) {
      throw new Error('Unix socket is not connected');
    }

    await new Promise<void>((resolve, reject) => {
      socket.write(serializeMessage(message), (error) => {
        if (error) {
          reject(error);
        } else {
          resolve();
        }
      });
    });
  }
}