import type { SmartNftables } from './nft.manager.js';
import type { INftFirewallRule, INftIPSetConfig } from './nft.types.js';
import {
  buildFirewallRule,
  buildIPSetCreate,
  buildIPSetAddElements,
  buildIPSetRemoveElements,
  buildIPSetDelete,
  buildIPSetMatchRule,
} from './nft.rulebuilder.firewall.js';

/**
 * Manages firewall (filter) rules, IP sets, and convenience methods.
 */
export class FirewallManager {
  constructor(private parent: SmartNftables) {}

  /**
   * Add a firewall rule (input/output/forward).
   */
  public async addRule(groupId: string, rule: INftFirewallRule): Promise<void> {
    await this.parent.ensureFilterChains();
    const commands = buildFirewallRule(this.parent.tableName, this.parent.family, rule);
    await this.parent.applyRuleGroup(`fw:${groupId}`, commands);
  }

  /**
   * Remove a firewall rule group.
   */
  public async removeRule(groupId: string): Promise<void> {
    await this.parent.removeRuleGroup(`fw:${groupId}`);
  }

  /**
   * Create a named IP set.
   */
  public async createIPSet(config: INftIPSetConfig): Promise<void> {
    await this.parent.ensureInitialized();
    const commands = buildIPSetCreate(this.parent.tableName, this.parent.family, config);
    await this.parent.applyRuleGroup(`ipset:${config.name}`, commands);
  }

  /**
   * Add elements to an existing IP set.
   */
  public async addToIPSet(setName: string, elements: string[]): Promise<void> {
    const commands = buildIPSetAddElements(this.parent.tableName, this.parent.family, setName, elements);
    if (commands.length > 0) {
      await this.parent.executor.execBatch(commands);
    }
  }

  /**
   * Remove elements from an IP set.
   */
  public async removeFromIPSet(setName: string, elements: string[]): Promise<void> {
    const commands = buildIPSetRemoveElements(this.parent.tableName, this.parent.family, setName, elements);
    if (commands.length > 0) {
      await this.parent.executor.execBatch(commands);
    }
  }

  /**
   * Delete an IP set entirely.
   */
  public async deleteIPSet(setName: string): Promise<void> {
    const commands = buildIPSetDelete(this.parent.tableName, this.parent.family, setName);
    await this.parent.executor.execBatch(commands);
    await this.parent.removeRuleGroup(`ipset:${setName}`);
  }

  /**
   * Convenience: block an IP or CIDR subnet.
   */
  public async blockIP(ip: string, options?: { direction?: 'input' | 'forward' }): Promise<void> {
    const direction = options?.direction ?? 'input';
    const safeId = ip.replace(/[/.]/g, '_');
    await this.addRule(`block-${safeId}`, {
      direction,
      action: 'drop',
      sourceIP: ip,
    });
  }

  /**
   * Convenience: allow only specific IPs on a port.
   * Adds accept rules for each IP, then a drop rule for everything else on that port.
   */
  public async allowOnlyIPs(
    groupId: string,
    ips: string[],
    port?: number,
    protocol?: 'tcp' | 'udp' | 'both',
  ): Promise<void> {
    await this.parent.ensureFilterChains();
    const commands: string[] = [];

    for (const ip of ips) {
      const acceptCmds = buildFirewallRule(this.parent.tableName, this.parent.family, {
        direction: 'input',
        action: 'accept',
        sourceIP: ip,
        destPort: port,
        protocol,
      });
      commands.push(...acceptCmds);
    }

    // Drop everything else on that port
    const dropCmds = buildFirewallRule(this.parent.tableName, this.parent.family, {
      direction: 'input',
      action: 'drop',
      destPort: port,
      protocol,
    });
    commands.push(...dropCmds);

    await this.parent.applyRuleGroup(`fw:allowonly:${groupId}`, commands);
  }

  /**
   * Convenience: enable stateful connection tracking.
   * Allows established+related, drops invalid.
   */
  public async enableStatefulTracking(chain?: 'input' | 'forward' | 'output'): Promise<void> {
    await this.parent.ensureFilterChains();
    const direction = chain ?? 'input';
    const commands: string[] = [];

    // Allow established and related connections
    const allowCmds = buildFirewallRule(this.parent.tableName, this.parent.family, {
      direction,
      action: 'accept',
      ctStates: ['established', 'related'],
    });
    commands.push(...allowCmds);

    // Drop invalid connections
    const dropCmds = buildFirewallRule(this.parent.tableName, this.parent.family, {
      direction,
      action: 'drop',
      ctStates: ['invalid'],
    });
    commands.push(...dropCmds);

    await this.parent.applyRuleGroup(`fw:stateful:${direction}`, commands);
  }
}