import type { McpInstance, McpServer } from '@prisma/client';
import type { IMcpInstanceRepository, IMcpServerRepository } from '../repositories/interfaces.js';
import type { McpOrchestrator } from './orchestrator.js';
import { NotFoundError } from './mcp-server.service.js';
import { InvalidStateError } from './instance.service.js';
import { sendViaSse } from './transport/sse-client.js';
import { sendViaStdio } from './transport/stdio-client.js';
import { PersistentStdioClient } from './transport/persistent-stdio.js';

/**
 * Build the spawn command for a runtime inside its runner container.
 * node → npx --prefer-offline -y <pkg>
 * python → uvx <pkg>
 */
export function buildRuntimeSpawnCmd(runtime: string, packageName: string): string[] {
  switch (runtime) {
    case 'python':
      return ['uvx', packageName];
    case 'node':
    default:
      return ['npx', '--prefer-offline', '-y', packageName];
  }
}

export interface McpProxyRequest {
  serverId: string;
  method: string;
  params?: Record<string, unknown> | undefined;
}

export interface McpProxyResponse {
  jsonrpc: '2.0';
  id: number;
  result?: unknown;
  error?: { code: number; message: string; data?: unknown };
}

/**
 * Parses a streamable-http SSE response body to extract the JSON-RPC payload.
 * Streamable-http returns `event: message\ndata: {...}\n\n` format.
 */
function parseStreamableResponse(body: string): McpProxyResponse {
  for (const line of body.split('\n')) {
    const trimmed = line.trim();
    if (trimmed.startsWith('data: ')) {
      return JSON.parse(trimmed.slice(6)) as McpProxyResponse;
    }
  }
  // If body is plain JSON (no SSE framing), parse directly
  return JSON.parse(body) as McpProxyResponse;
}

export class McpProxyService {
  /** Session IDs per server for streamable-http protocol */
  private sessions = new Map<string, string>();
  /** Persistent STDIO connections keyed by containerId */
  private stdioClients = new Map<string, PersistentStdioClient>();

  constructor(
    private readonly instanceRepo: IMcpInstanceRepository,
    private readonly serverRepo: IMcpServerRepository,
    private readonly orchestrator?: McpOrchestrator,
  ) {}

  /** Clean up all persistent connections (call on shutdown). */
  closeAll(): void {
    for (const [, client] of this.stdioClients) {
      client.close();
    }
    this.stdioClients.clear();
  }

  /** Remove persistent connection for a container (call when instance stops). */
  removeClient(containerId: string): void {
    const client = this.stdioClients.get(containerId);
    if (client) {
      client.close();
      this.stdioClients.delete(containerId);
    }
  }

  async execute(request: McpProxyRequest): Promise<McpProxyResponse> {
    const server = await this.serverRepo.findById(request.serverId);
    if (!server) {
      throw new NotFoundError(`Server '${request.serverId}' not found`);
    }

    // External server: proxy directly to externalUrl
    if (server.externalUrl) {
      return this.sendToExternal(server, request.method, request.params);
    }

    // Managed server: find running instance and dispatch by transport
    const instances = await this.instanceRepo.findAll(request.serverId);
    const running = instances.find((i) => i.status === 'RUNNING');

    if (!running) {
      throw new NotFoundError(`No running instance found for server '${request.serverId}'`);
    }

    return this.sendToManaged(server, running, request.method, request.params);
  }

  /**
   * Send to an external MCP server. Dispatches based on transport type.
   */
  private async sendToExternal(
    server: McpServer,
    method: string,
    params?: Record<string, unknown>,
  ): Promise<McpProxyResponse> {
    const url = server.externalUrl as string;

    if (server.transport === 'SSE') {
      return sendViaSse(url, method, params);
    }

    // STREAMABLE_HTTP (default for external)
    return this.sendStreamableHttp(server.id, url, method, params);
  }

  /**
   * Send to a managed (containerized) MCP server. Dispatches based on transport type.
   */
  private async sendToManaged(
    server: McpServer,
    instance: McpInstance,
    method: string,
    params?: Record<string, unknown>,
  ): Promise<McpProxyResponse> {
    const transport = server.transport as string;

    // STDIO: use persistent connection (falls back to one-shot on error)
    if (transport === 'STDIO') {
      if (!this.orchestrator) {
        throw new InvalidStateError('Orchestrator required for STDIO transport');
      }
      if (!instance.containerId) {
        throw new InvalidStateError(`Instance '${instance.id}' has no container ID`);
      }
      const packageName = server.packageName as string | null;
      const command = server.command as string[] | null;
      if (!packageName && (!command || command.length === 0)) {
        throw new InvalidStateError(`Server '${server.id}' has no packageName or command for STDIO transport`);
      }

      // Build the spawn command based on runtime
      const runtime = (server.runtime as string | null) ?? 'node';
      const spawnCmd = command && command.length > 0
        ? command
        : buildRuntimeSpawnCmd(runtime, packageName!);

      // Try persistent connection first
      try {
        return await this.sendViaPersistentStdio(instance.containerId, spawnCmd, method, params);
      } catch {
        // Persistent failed — fall back to one-shot
        this.removeClient(instance.containerId);
        return sendViaStdio(this.orchestrator, instance.containerId, packageName, method, params, 120_000, command, runtime);
      }
    }

    // SSE or STREAMABLE_HTTP: need a base URL
    const baseUrl = await this.resolveBaseUrl(instance, server);

    if (transport === 'SSE') {
      return sendViaSse(baseUrl, method, params);
    }

    // STREAMABLE_HTTP (default)
    return this.sendStreamableHttp(server.id, baseUrl, method, params);
  }

  /**
   * Send via a persistent STDIO connection (reused across calls).
   */
  private async sendViaPersistentStdio(
    containerId: string,
    command: string[],
    method: string,
    params?: Record<string, unknown>,
  ): Promise<McpProxyResponse> {
    let client = this.stdioClients.get(containerId);
    if (!client) {
      client = new PersistentStdioClient(this.orchestrator!, containerId, command);
      this.stdioClients.set(containerId, client);
    }
    return client.send(method, params);
  }

  /**
   * Resolve the base URL for an HTTP-based managed server.
   * Prefers container internal IP on Docker network, falls back to localhost:port.
   */
  private async resolveBaseUrl(instance: McpInstance, server: McpServer): Promise<string> {
    const containerPort = (server.containerPort as number | null) ?? 3000;

    if (this.orchestrator && instance.containerId) {
      try {
        const containerInfo = await this.orchestrator.inspectContainer(instance.containerId);
        if (containerInfo.ip) {
          return `http://${containerInfo.ip}:${containerPort}`;
        }
      } catch {
        // Fall through to localhost
      }
    }

    if (instance.port !== null && instance.port !== undefined) {
      return `http://localhost:${instance.port}`;
    }

    throw new InvalidStateError(
      `Cannot resolve URL for instance '${instance.id}': no container IP or host port`,
    );
  }

  /**
   * Send via streamable-http protocol with session management.
   */
  private async sendStreamableHttp(
    serverId: string,
    url: string,
    method: string,
    params?: Record<string, unknown>,
  ): Promise<McpProxyResponse> {
    // Ensure we have a session (initialize on first call)
    if (!this.sessions.has(serverId)) {
      await this.initSession(serverId, url);
    }

    const sessionId = this.sessions.get(serverId);

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

    const headers: Record<string, string> = {
      'Content-Type': 'application/json',
      'Accept': 'application/json, text/event-stream',
    };
    if (sessionId) {
      headers['Mcp-Session-Id'] = sessionId;
    }

    const response = await fetch(url, {
      method: 'POST',
      headers,
      body: JSON.stringify(body),
    });

    if (!response.ok) {
      // Session expired? Clear and retry once
      if (response.status === 400 || response.status === 404) {
        this.sessions.delete(serverId);
        return this.sendStreamableHttp(serverId, url, method, params);
      }
      return {
        jsonrpc: '2.0',
        id: 1,
        error: {
          code: -32000,
          message: `MCP server returned HTTP ${response.status}: ${response.statusText}`,
        },
      };
    }

    const text = await response.text();
    return parseStreamableResponse(text);
  }

  /**
   * Initialize a streamable-http session with a server.
   */
  private async initSession(serverId: string, url: string): Promise<void> {
    const initBody = {
      jsonrpc: '2.0',
      id: 1,
      method: 'initialize',
      params: {
        protocolVersion: '2025-03-26',
        capabilities: {},
        clientInfo: { name: 'mcpctl', version: '0.0.1' },
      },
    };

    const response = await fetch(url, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        'Accept': 'application/json, text/event-stream',
      },
      body: JSON.stringify(initBody),
    });

    if (!response.ok) {
      throw new Error(`Failed to initialize session: HTTP ${response.status}`);
    }

    const sessionId = response.headers.get('mcp-session-id');
    if (sessionId) {
      this.sessions.set(serverId, sessionId);
    }

    // Send notifications/initialized
    const headers: Record<string, string> = {
      'Content-Type': 'application/json',
      'Accept': 'application/json, text/event-stream',
    };
    if (sessionId) {
      headers['Mcp-Session-Id'] = sessionId;
    }

    await fetch(url, {
      method: 'POST',
      headers,
      body: JSON.stringify({ jsonrpc: '2.0', method: 'notifications/initialized' }),
    });
  }
}
