// firewall.c - Firewall module for utun
#define _POSIX_C_SOURCE 200809L
#include "firewall.h"
#include <stdlib.h>
#include <string.h>
#include "../lib/debug_config.h"
#include "../lib/mem.h"

static int compare_rules(const void *a, const void *b) {
    const struct CFG_FIREWALL_RULE *ra = (const struct CFG_FIREWALL_RULE *)a;
    const struct CFG_FIREWALL_RULE *rb = (const struct CFG_FIREWALL_RULE *)b;
    if (ra->ip != rb->ip) {
        return (ra->ip > rb->ip) ? 1 : -1;
    }
    return (ra->port > rb->port) ? 1 : -1;
}

int fw_init(struct firewall_ctx *ctx) {
    if (!ctx) return -1;
    memset(ctx, 0, sizeof(struct firewall_ctx));
    return 0;
}

int fw_load_rules(struct firewall_ctx *ctx, const struct global_config *cfg) {
    if (!ctx || !cfg) return -1;

    fw_free(ctx);
    fw_init(ctx);

    ctx->bypass_all = cfg->firewall_bypass_all;

    if (cfg->firewall_rule_count > 0) {
        ctx->rules = u_calloc(cfg->firewall_rule_count, sizeof(struct CFG_FIREWALL_RULE));
        if (!ctx->rules) {
            DEBUG_ERROR(DEBUG_CATEGORY_MEMORY, "fw_load_rules: failed to allocate rules");
            return -1;
        }
        memcpy(ctx->rules, cfg->firewall_rules, 
               cfg->firewall_rule_count * sizeof(struct CFG_FIREWALL_RULE));
        ctx->count = cfg->firewall_rule_count;

        // Sort rules by IP (and port)
        qsort(ctx->rules, ctx->count, sizeof(struct CFG_FIREWALL_RULE), compare_rules);
        
        DEBUG_INFO(DEBUG_CATEGORY_CONFIG, "Loaded %d firewall rules (sorted)", ctx->count);
    }

    if (ctx->bypass_all) {
        DEBUG_INFO(DEBUG_CATEGORY_CONFIG, "Firewall bypass_all mode enabled");
    }

    return 0;
}

int fw_check(const struct firewall_ctx *ctx, uint32_t ip, uint16_t port) {
    if (!ctx) return 0;
    if (ctx->bypass_all) return 1;

    // Binary search for IP
    int low = 0;
    int high = ctx->count - 1;
    while (low <= high) {
        int mid = (low + high) / 2;
        if (ctx->rules[mid].ip == ip) {
            // Found matching IP, check port
            int i = mid;
            // Check exact match or any-port rule
            while (i >= 0 && ctx->rules[i].ip == ip) {
                if (ctx->rules[i].port == 0 || ctx->rules[i].port == port) {
                    return 1;
                }
                i--;
            }
            i = mid + 1;
            while (i < ctx->count && ctx->rules[i].ip == ip) {
                if (ctx->rules[i].port == 0 || ctx->rules[i].port == port) {
                    return 1;
                }
                i++;
            }
            return 0;
        } else if (ctx->rules[mid].ip < ip) {
            low = mid + 1;
        } else {
            high = mid - 1;
        }
    }
    return 0;
}

void fw_free(struct firewall_ctx *ctx) {
    if (!ctx) return;
    u_free(ctx->rules);
    memset(ctx, 0, sizeof(struct firewall_ctx));
}