import qs from 'qs';
import yargs from 'yargs';
import { hideBin } from 'yargs/helpers';
import z from 'zod';

export type CLIOptions = McpOptions & {
  transport: 'stdio' | 'http';
  port: number | undefined;
  socket: string | undefined;
};

export type McpOptions = {
  includeDocsTools?: boolean | undefined;
  includeGetAgenticHelpTool?: boolean | undefined;
};

export function parseCLIOptions(): CLIOptions {
  const opts = yargs(hideBin(process.argv))
    .option('tools', {
      type: 'string',
      array: true,
      choices: ['code', 'docs', 'help'],
      description: 'Use dynamic tools or all tools',
    })
    .option('no-tools', {
      type: 'string',
      array: true,
      choices: ['code', 'docs', 'help'],
      description: 'Do not use any dynamic or all tools',
    })
    .option('transport', {
      type: 'string',
      choices: ['stdio', 'http'],
      default: 'stdio',
      description: 'What transport to use; stdio for local servers or http for remote servers',
    })
    .option('port', {
      type: 'number',
      description: 'Port to serve on if using http transport',
    })
    .option('socket', {
      type: 'string',
      description: 'Unix socket to serve on if using http transport',
    })
    .help();

  const argv = opts.parseSync();

  const shouldIncludeToolType = (toolType: 'code' | 'docs' | 'help') =>
    argv.noTools?.includes(toolType) ? false
    : argv.tools?.includes(toolType) ? true
    : undefined;

  const includeDocsTools = shouldIncludeToolType('docs');
  const includeGetAgenticHelpTool = shouldIncludeToolType('help');

  const transport = argv.transport as 'stdio' | 'http';

  return {
    ...(includeDocsTools !== undefined && { includeDocsTools }),
    ...(includeGetAgenticHelpTool !== undefined && { includeGetAgenticHelpTool }),
    transport,
    port: argv.port,
    socket: argv.socket,
  };
}

const coerceArray = <T extends z.ZodTypeAny>(zodType: T) =>
  z.preprocess(
    (val) =>
      Array.isArray(val) ? val
      : val ? [val]
      : val,
    z.array(zodType).optional(),
  );

const QueryOptions = z.object({
  tools: coerceArray(z.enum(['code', 'docs', 'help'])).describe('Specify which MCP tools to use'),
  no_tools: coerceArray(z.enum(['code', 'docs', 'help'])).describe('Specify which MCP tools to not use.'),
  tool: coerceArray(z.string()).describe('Include tools matching the specified names'),
});

export function parseQueryOptions(defaultOptions: McpOptions, query: unknown): McpOptions {
  const queryObject = typeof query === 'string' ? qs.parse(query) : query;
  const queryOptions = QueryOptions.parse(queryObject);

  let docsTools: boolean | undefined =
    queryOptions.no_tools && queryOptions.no_tools?.includes('docs') ? false
    : queryOptions.tools?.includes('docs') ? true
    : defaultOptions.includeDocsTools;

  let getAgenticHelpTool: boolean | undefined =
    queryOptions.no_tools && queryOptions.no_tools?.includes('help') ? false
    : queryOptions.tools?.includes('help') && defaultOptions.includeGetAgenticHelpTool ? true
    : defaultOptions.includeGetAgenticHelpTool;

  return {
    ...(docsTools !== undefined && { includeDocsTools: docsTools }),
    ...(getAgenticHelpTool !== undefined && { includeGetAgenticHelpTool: getAgenticHelpTool }),
  };
}