import { spawn, type ChildProcess } from 'node:child_process';
import { createInterface } from 'node:readline';
import type { UpstreamConnection, JsonRpcRequest, JsonRpcResponse, JsonRpcNotification, UpstreamConfig } from '../types.js';

/**
 * Connects to an MCP server over STDIO (spawn a process, write JSON-RPC to stdin, read from stdout).
 */
export class StdioUpstream implements UpstreamConnection {
  readonly name: string;
  private process: ChildProcess | null = null;
  private pendingRequests = new Map<string | number, {
    resolve: (res: JsonRpcResponse) => void;
    reject: (err: Error) => void;
  }>();
  private alive = false;
  private notificationHandlers: Array<(notification: JsonRpcNotification) => void> = [];

  constructor(private config: UpstreamConfig) {
    this.name = config.name;
  }

  async start(): Promise<void> {
    if (!this.config.command) {
      throw new Error(`STDIO upstream '${this.name}' has no command configured`);
    }

    this.process = spawn(this.config.command, this.config.args ?? [], {
      stdio: ['pipe', 'pipe', 'pipe'],
      env: { ...process.env, ...this.config.env },
    });

    this.alive = true;

    this.process.on('exit', () => {
      this.alive = false;
      // Reject any pending requests
      for (const [, pending] of this.pendingRequests) {
        pending.reject(new Error(`Upstream '${this.name}' process exited`));
      }
      this.pendingRequests.clear();
    });

    if (this.process.stdout) {
      const rl = createInterface({ input: this.process.stdout });
      rl.on('line', (line) => {
        try {
          const msg = JSON.parse(line) as Record<string, unknown>;
          if ('id' in msg && msg.id !== undefined) {
            // Response to a pending request
            const pending = this.pendingRequests.get(msg.id as string | number);
            if (pending) {
              this.pendingRequests.delete(msg.id as string | number);
              pending.resolve(msg as unknown as JsonRpcResponse);
            }
          } else if ('method' in msg) {
            // Notification from upstream
            const notification: JsonRpcNotification = {
              jsonrpc: '2.0',
              method: msg.method as string,
            };
            if (msg.params) {
              notification.params = msg.params as Record<string, unknown>;
            }
            for (const handler of this.notificationHandlers) {
              handler(notification);
            }
          }
        } catch {
          // Skip non-JSON lines
        }
      });
    }
  }

  async send(request: JsonRpcRequest): Promise<JsonRpcResponse> {
    if (!this.process?.stdin || !this.alive) {
      throw new Error(`Upstream '${this.name}' is not connected`);
    }

    return new Promise((resolve, reject) => {
      this.pendingRequests.set(request.id, { resolve, reject });

      const timeout = setTimeout(() => {
        this.pendingRequests.delete(request.id);
        reject(new Error(`Request to '${this.name}' timed out`));
      }, 30000);

      this.pendingRequests.set(request.id, {
        resolve: (res) => {
          clearTimeout(timeout);
          resolve(res);
        },
        reject: (err) => {
          clearTimeout(timeout);
          reject(err);
        },
      });

      this.process!.stdin!.write(JSON.stringify(request) + '\n');
    });
  }

  async close(): Promise<void> {
    if (this.process) {
      this.process.kill('SIGTERM');
      this.alive = false;
      this.process = null;
    }
  }

  isAlive(): boolean {
    return this.alive;
  }

  onNotification(handler: (notification: JsonRpcNotification) => void): void {
    this.notificationHandlers.push(handler);
  }
}
