import type { TNftFamily, INftRateLimitRule, INftConnectionRateRule } from './nft.types.js';

/**
 * Expand a protocol spec into concrete protocol strings.
 */
function expandProtocols(protocol?: 'tcp' | 'udp' | 'both'): string[] {
  switch (protocol ?? 'tcp') {
    case 'tcp': return ['tcp'];
    case 'udp': return ['udp'];
    case 'both': return ['tcp', 'udp'];
  }
}

/**
 * Build a rate limit rule.
 * Packets exceeding the rate are subjected to the specified action (default: drop).
 */
export function buildRateLimitRule(
  tableName: string,
  family: TNftFamily,
  rule: INftRateLimitRule,
): string[] {
  const protocols = expandProtocols(rule.protocol);
  const chain = rule.chain ?? 'input';
  const action = rule.action ?? 'drop';
  const commands: string[] = [];

  for (const proto of protocols) {
    const portMatch = rule.port != null ? ` ${proto} dport ${rule.port}` : '';
    const burstClause = rule.burst != null ? ` burst ${rule.burst} packets` : '';

    if (rule.perSourceIP) {
      // Per-IP rate limiting using nft meters
      const meterName = `meter_${proto}_${rule.port ?? 'all'}`;
      commands.push(
        `nft add rule ${family} ${tableName} ${chain}${portMatch} meter ${meterName} { ip saddr limit rate over ${rule.rate}${burstClause} } ${action}`
      );
    } else {
      // Global rate limiting
      commands.push(
        `nft add rule ${family} ${tableName} ${chain}${portMatch} limit rate over ${rule.rate}${burstClause} ${action}`
      );
    }
  }

  return commands;
}

/**
 * Build a per-IP rate limit rule using nft meters.
 * Convenience wrapper around buildRateLimitRule with perSourceIP=true.
 */
export function buildPerIpRateLimitRule(
  tableName: string,
  family: TNftFamily,
  rule: Omit<INftRateLimitRule, 'perSourceIP'>,
): string[] {
  return buildRateLimitRule(tableName, family, { ...rule, perSourceIP: true });
}

/**
 * Build a new-connection rate limit rule.
 * Limits the rate of new connections (ct state new) on the given port.
 */
export function buildConnectionRateRule(
  tableName: string,
  family: TNftFamily,
  rule: INftConnectionRateRule,
): string[] {
  const protocols = expandProtocols(rule.protocol);
  const commands: string[] = [];

  for (const proto of protocols) {
    const portMatch = rule.port != null ? ` ${proto} dport ${rule.port}` : '';

    if (rule.perSourceIP) {
      const meterName = `connrate_${proto}_${rule.port ?? 'all'}`;
      commands.push(
        `nft add rule ${family} ${tableName} input ct state new${portMatch} meter ${meterName} { ip saddr limit rate over ${rule.rate} } drop`
      );
    } else {
      commands.push(
        `nft add rule ${family} ${tableName} input ct state new${portMatch} limit rate over ${rule.rate} drop`
      );
    }
  }

  return commands;
}