import { Client } from '@modelcontextprotocol/sdk/client/index.js';
import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
import { checkSSRFStrict, pinnedFetch } from './ssrf-strict.js';
import type { McpServerRecord } from './types.js';
import { logger } from '../logger.js';

/**
 * Build the auth headers sent to an MCP server.
 *
 * api_key servers may declare a custom auth header (e.g. NocoDB's `xc-mcp-token`),
 * in which case the token is sent verbatim under that header name. Otherwise
 * (and always for oauth) the standard `Authorization: Bearer <token>` is used.
 */
export function buildAuthHeaders(
  server: Pick<McpServerRecord, 'authKind' | 'authHeaderName'>,
  accessToken: string,
): Record<string, string> {
  if (server.authKind === 'api_key' && server.authHeaderName) {
    return { [server.authHeaderName]: accessToken };
  }
  return { Authorization: `Bearer ${accessToken}` };
}

export interface ClientFactoryOptions {
  insecureLocalTestMode?: boolean;
  /**
   * When true, skip the SSRF strict-check for private/loopback addresses.
   * Controlled by config.yaml `mcp.allow_private_addresses`.
   */
  allowPrivateAddresses?: boolean;
  callTimeoutMs: number;
}

export async function createMcpClient(
  server: Pick<McpServerRecord, 'id' | 'url' | 'authKind' | 'authHeaderName'>,
  accessToken: string,
  opts: ClientFactoryOptions,
): Promise<{ client: Client; close: () => Promise<void> }> {
  let fetcher: typeof fetch;
  if (opts.insecureLocalTestMode || opts.allowPrivateAddresses) {
    fetcher = fetch;
  } else {
    const ssrf = await checkSSRFStrict(server.url);
    if (!ssrf.ok) throw new Error(`SSRF check failed for server '${server.id}': ${ssrf.reason}`);
    fetcher = ((url: string, init?: RequestInit) =>
      pinnedFetch(url, {
        ...(init ?? {}),
        pinnedIp: ssrf.pinnedIp,
        family: ssrf.family,
      })) as typeof fetch;
  }

  const transport = new StreamableHTTPClientTransport(new URL(server.url), {
    fetch: fetcher,
    requestInit: {
      headers: buildAuthHeaders(server, accessToken),
    },
  });

  const client = new Client(
    { name: 'maestro', version: '0.1.0' },
    { capabilities: {} },
  );

  await client.connect(transport);
  logger.debug(`[mcp:client] connected server=${server.id}`);

  return {
    client,
    close: async () => {
      try {
        await client.close();
      } catch (err) {
        logger.warn(`[mcp:client] close failed server=${server.id}: ${(err as Error).message}`);
      }
    },
  };
}
