import 'dart:io';

class TrafficControlService {
  static Future<void> setSpeedLimit(String peerIP, int speedKbps) async {
    final mark = _getMarkForIP(peerIP);

    await _runIptablesCommand(['-I', 'FORWARD', '-s', peerIP, '-j', 'MARK', '--set-mark', mark.toString()]);
    await _runIptablesCommand(['-I', 'FORWARD', '-d', peerIP, '-j', 'MARK', '--set-mark', mark.toString()]);
    
    await _runTcCommand(['class', 'add', 'dev', 'wg0', 'parent', '1:1', 'classid', '1:$mark', 'htb', 'rate', '${speedKbps}kbit', 'ceil', '${speedKbps}kbit']);
    await _runTcCommand(['filter', 'add', 'dev', 'wg0', 'protocol', 'ip', 'parent', '1:', 'prio', '1', 'handle', mark.toString(), 'fw', 'flowid', '1:$mark']);
  }

  static Future<void> setDataCap(String peerIP, int dataCapMB) async {
    final quotaBytes = dataCapMB * 1024 * 1024;
    
    await _runIptablesCommand(['-I', 'FORWARD', '-s', peerIP, '-m', 'quota', '--quota', quotaBytes.toString(), '-j', 'ACCEPT']);
    await _runIptablesCommand(['-I', 'FORWARD', '-d', peerIP, '-m', 'quota', '--quota', quotaBytes.toString(), '-j', 'ACCEPT']);
    await _runIptablesCommand(['-A', 'FORWARD', '-s', peerIP, '-j', 'DROP']);
    await _runIptablesCommand(['-A', 'FORWARD', '-d', peerIP, '-j', 'DROP']);
  }

  static int _getMarkForIP(String ip) {
    return ip.split('.').last.hashCode % 65535 + 1;
  }

  static Future<void> _runIptablesCommand(List<String> args) async {
    final result = await Process.run('iptables', args);
    if (result.exitCode != 0) {
      throw Exception('iptables command failed: ${result.stderr}');
    }
  }

  static Future<void> _runTcCommand(List<String> args) async {
    final result = await Process.run('tc', args);
    if (result.exitCode != 0) {
      throw Exception('tc command failed: ${result.stderr}');
    }
  }
}