import type { McpOrchestrator, InteractiveExec } from '../orchestrator.js';
import type { McpProxyResponse } from '../mcp-proxy-service.js';

/**
 * Persistent STDIO connection to an MCP server running inside a Docker container.
 *
 * Instead of cold-starting a new process per call (docker exec one-shot), this keeps
 * a long-running `docker exec -i <cmd>` session alive. The MCP init handshake runs
 * once, then tool calls are multiplexed over the same stdin/stdout pipe.
 *
 * Falls back gracefully: if the process dies, the next call will reconnect.
 */
export class PersistentStdioClient {
  private exec: InteractiveExec | null = null;
  private buffer = '';
  private nextId = 1;
  private initialized = false;
  private connecting: Promise<void> | null = null;
  private pendingRequests = new Map<number, {
    resolve: (res: McpProxyResponse) => void;
    reject: (err: Error) => void;
    timer: ReturnType<typeof setTimeout>;
  }>();

  constructor(
    private readonly orchestrator: McpOrchestrator,
    private readonly containerId: string,
    private readonly command: string[],
    private readonly timeoutMs = 120_000,
  ) {}

  /**
   * Send a JSON-RPC request and wait for the matching response.
   */
  async send(method: string, params?: Record<string, unknown>): Promise<McpProxyResponse> {
    await this.ensureReady();

    const id = this.nextId++;
    const request: Record<string, unknown> = { jsonrpc: '2.0', id, method };
    if (params !== undefined) {
      request.params = params;
    }

    return new Promise<McpProxyResponse>((resolve, reject) => {
      const timer = setTimeout(() => {
        this.pendingRequests.delete(id);
        reject(new Error(`Request timed out after ${this.timeoutMs}ms`));
      }, this.timeoutMs);

      this.pendingRequests.set(id, { resolve, reject, timer });
      this.write(request);
    });
  }

  /** Shut down the persistent connection. */
  close(): void {
    if (this.exec) {
      this.exec.close();
      this.exec = null;
    }
    this.initialized = false;
    this.connecting = null;
    for (const [, pending] of this.pendingRequests) {
      clearTimeout(pending.timer);
      pending.reject(new Error('Connection closed'));
    }
    this.pendingRequests.clear();
  }

  get isConnected(): boolean {
    return this.initialized && this.exec !== null;
  }

  // ── internals ──

  private async ensureReady(): Promise<void> {
    if (this.initialized && this.exec) return;
    if (this.connecting) {
      await this.connecting;
      return;
    }
    this.connecting = this.connect();
    try {
      await this.connecting;
    } finally {
      this.connecting = null;
    }
  }

  private async connect(): Promise<void> {
    this.close();

    if (!this.orchestrator.execInteractive) {
      throw new Error('Orchestrator does not support interactive exec');
    }

    const exec = await this.orchestrator.execInteractive(this.containerId, this.command);
    this.exec = exec;
    this.buffer = '';

    // Parse JSON-RPC responses line by line from stdout
    exec.stdout.on('data', (chunk: Buffer) => {
      this.buffer += chunk.toString('utf-8');
      this.processBuffer();
    });

    exec.stdout.on('end', () => {
      this.initialized = false;
      this.exec = null;
      for (const [, pending] of this.pendingRequests) {
        clearTimeout(pending.timer);
        pending.reject(new Error('STDIO process exited'));
      }
      this.pendingRequests.clear();
    });

    // Run MCP init handshake
    const initId = this.nextId++;
    const initPromise = new Promise<void>((resolve, reject) => {
      const timer = setTimeout(() => {
        this.pendingRequests.delete(initId);
        reject(new Error('MCP init handshake timed out'));
      }, 30_000);

      this.pendingRequests.set(initId, {
        resolve: () => {
          clearTimeout(timer);
          resolve();
        },
        reject: (err) => {
          clearTimeout(timer);
          reject(err);
        },
        timer,
      });
    });

    this.write({
      jsonrpc: '2.0',
      id: initId,
      method: 'initialize',
      params: {
        protocolVersion: '2024-11-05',
        capabilities: {},
        clientInfo: { name: 'mcpctl-proxy', version: '0.0.1' },
      },
    });

    await initPromise;

    // Send initialized notification (no response expected)
    this.write({ jsonrpc: '2.0', method: 'notifications/initialized' });

    // Small delay to let the server process the notification
    await new Promise((r) => setTimeout(r, 100));

    this.initialized = true;
  }

  private write(msg: Record<string, unknown>): void {
    if (!this.exec) throw new Error('Not connected');
    this.exec.write(JSON.stringify(msg) + '\n');
  }

  private processBuffer(): void {
    const lines = this.buffer.split('\n');
    this.buffer = lines.pop() ?? '';

    for (const line of lines) {
      const trimmed = line.trim();
      if (!trimmed) continue;
      try {
        const msg = JSON.parse(trimmed) as Record<string, unknown>;
        if ('id' in msg && msg.id !== undefined) {
          const pending = this.pendingRequests.get(msg.id as number);
          if (pending) {
            this.pendingRequests.delete(msg.id as number);
            clearTimeout(pending.timer);
            pending.resolve(msg as unknown as McpProxyResponse);
          }
        }
        // Notifications from server are ignored (not needed for proxy)
      } catch {
        // Skip non-JSON lines
      }
    }
  }
}
