#define _GNU_SOURCE
#include "ip_ban.h"
#include "debug.h"
#include "config.h"
#include <string.h>
#include <stdlib.h>
#include <pthread.h>

// ============================================================
// IP Auth Failure Ban System
//
// Fixed-size open-addressing hash table. No malloc after init.
// Thread-safe via a single mutex (low contention — only called
// at connection open/close, not in the hot event path).
// ============================================================

#define IP_BAN_EMPTY    0   // slot is unused
#define IP_BAN_ACTIVE   1   // slot has an entry

typedef struct {
    int state;                  // IP_BAN_EMPTY or IP_BAN_ACTIVE
    char ip[46];                // IPv4 or IPv6 string
    int failure_count;          // total failures in current window
    time_t first_failure;       // start of current failure window
    time_t banned_until;        // 0 = not banned; >0 = banned until this time
    int ban_count;              // how many times this IP has been banned (for backoff)
} ip_ban_entry_t;

static ip_ban_entry_t g_ban_table[IP_BAN_TABLE_SIZE];
static pthread_mutex_t g_ban_mutex = PTHREAD_MUTEX_INITIALIZER;
static int g_initialized = 0;

// Simple FNV-1a hash for IP strings
static unsigned int ip_hash(const char* ip) {
    unsigned int hash = 2166136261u;
    while (*ip) {
        hash ^= (unsigned char)*ip++;
        hash *= 16777619u;
    }
    return hash % IP_BAN_TABLE_SIZE;
}

// Find slot for IP (open addressing with linear probing)
// Returns index of existing entry or first empty slot, -1 if table full
static int find_slot(const char* ip) {
    unsigned int start = ip_hash(ip);
    for (unsigned int i = 0; i < IP_BAN_TABLE_SIZE; i++) {
        unsigned int idx = (start + i) % IP_BAN_TABLE_SIZE;
        if (g_ban_table[idx].state == IP_BAN_EMPTY) {
            return (int)idx;  // empty slot — can insert here
        }
        if (strcmp(g_ban_table[idx].ip, ip) == 0) {
            return (int)idx;  // found existing entry
        }
    }
    return -1;  // table full (shouldn't happen with 4096 slots)
}

void ip_ban_init(void) {
    pthread_mutex_lock(&g_ban_mutex);
    memset(g_ban_table, 0, sizeof(g_ban_table));
    g_initialized = 1;
    pthread_mutex_unlock(&g_ban_mutex);
    DEBUG_LOG("IP ban table initialized (%d slots)", IP_BAN_TABLE_SIZE);
}

int ip_ban_is_banned(const char* ip) {
    if (!ip || !g_initialized) return 0;

    // Check if feature is enabled
    if (!get_config_bool("auth_fail_ban_enabled", 1)) return 0;

    pthread_mutex_lock(&g_ban_mutex);

    int idx = find_slot(ip);
    if (idx < 0 || g_ban_table[idx].state == IP_BAN_EMPTY) {
        pthread_mutex_unlock(&g_ban_mutex);
        return 0;  // no entry — not banned
    }

    ip_ban_entry_t* entry = &g_ban_table[idx];
    time_t now = time(NULL);

    if (entry->banned_until > 0 && now < entry->banned_until) {
        pthread_mutex_unlock(&g_ban_mutex);
        DEBUG_TRACE("IP %s is banned for %ld more seconds", ip, entry->banned_until - now);
        return 1;  // still banned
    }

    // Ban expired — clear it but keep the entry for failure tracking
    if (entry->banned_until > 0 && now >= entry->banned_until) {
        entry->banned_until = 0;
        entry->failure_count = 0;
        entry->first_failure = 0;
    }

    pthread_mutex_unlock(&g_ban_mutex);
    return 0;
}

void ip_ban_record_failure(const char* ip) {
    if (!ip || !g_initialized) return;
    if (!get_config_bool("auth_fail_ban_enabled", 1)) return;

    int threshold = get_config_int("auth_fail_ban_threshold", 3);
    int window_sec = get_config_int("auth_fail_window_sec", 60);
    int ban_duration = get_config_int("auth_fail_ban_duration_sec", 300);

    pthread_mutex_lock(&g_ban_mutex);

    int idx = find_slot(ip);
    if (idx < 0) {
        // Table full — can't track this IP, just log and return
        DEBUG_WARN("IP ban table full, cannot track %s", ip);
        pthread_mutex_unlock(&g_ban_mutex);
        return;
    }

    ip_ban_entry_t* entry = &g_ban_table[idx];
    time_t now = time(NULL);

    if (entry->state == IP_BAN_EMPTY) {
        // New entry
        entry->state = IP_BAN_ACTIVE;
        strncpy(entry->ip, ip, sizeof(entry->ip) - 1);
        entry->ip[sizeof(entry->ip) - 1] = '\0';
        entry->failure_count = 0;
        entry->first_failure = now;
        entry->banned_until = 0;
        entry->ban_count = 0;
    }

    // Reset window if it's expired
    if (entry->first_failure > 0 && (now - entry->first_failure) > window_sec) {
        entry->failure_count = 0;
        entry->first_failure = now;
    }

    entry->failure_count++;

    DEBUG_TRACE("IP %s auth failure count: %d/%d", ip, entry->failure_count, threshold);

    if (entry->failure_count >= threshold) {
        // Apply exponential backoff: ban_duration * 2^ban_count, capped at 24 hours
        int duration = ban_duration;
        for (int i = 0; i < entry->ban_count && duration < 86400; i++) {
            duration *= 2;
        }
        if (duration > 86400) duration = 86400;

        entry->banned_until = now + duration;
        entry->ban_count++;
        entry->failure_count = 0;  // reset counter after ban
        entry->first_failure = 0;

        DEBUG_WARN("IP %s banned for %d seconds (ban #%d) after %d auth failures",
                   ip, duration, entry->ban_count, threshold);
    }

    pthread_mutex_unlock(&g_ban_mutex);
}

void ip_ban_record_success(const char* ip) {
    if (!ip || !g_initialized) return;

    pthread_mutex_lock(&g_ban_mutex);

    int idx = find_slot(ip);
    if (idx >= 0 && g_ban_table[idx].state == IP_BAN_ACTIVE) {
        // Clear failure count on successful auth — reward good behavior
        g_ban_table[idx].failure_count = 0;
        g_ban_table[idx].first_failure = 0;
        // Note: we keep ban_count so backoff persists across sessions
        DEBUG_TRACE("IP %s authenticated successfully — failure count cleared", ip);
    }

    pthread_mutex_unlock(&g_ban_mutex);
}

void ip_ban_cleanup(void) {
    if (!g_initialized) return;

    pthread_mutex_lock(&g_ban_mutex);

    time_t now = time(NULL);
    int window_sec = get_config_int("auth_fail_window_sec", 60);
    int cleaned = 0;

    for (int i = 0; i < IP_BAN_TABLE_SIZE; i++) {
        if (g_ban_table[i].state != IP_BAN_ACTIVE) continue;

        ip_ban_entry_t* entry = &g_ban_table[i];

        // Clear entries where ban has expired AND failure window has expired AND no recent activity
        int ban_expired = (entry->banned_until == 0 || now >= entry->banned_until);
        int window_expired = (entry->first_failure == 0 || (now - entry->first_failure) > window_sec * 10);

        if (ban_expired && window_expired && entry->failure_count == 0) {
            memset(entry, 0, sizeof(ip_ban_entry_t));
            cleaned++;
        }
    }

    if (cleaned > 0) {
        DEBUG_TRACE("IP ban cleanup: freed %d stale entries", cleaned);
    }

    pthread_mutex_unlock(&g_ban_mutex);
}

int ip_ban_get_banned_count(void) {
    if (!g_initialized) return 0;

    pthread_mutex_lock(&g_ban_mutex);
    time_t now = time(NULL);
    int count = 0;
    for (int i = 0; i < IP_BAN_TABLE_SIZE; i++) {
        if (g_ban_table[i].state == IP_BAN_ACTIVE &&
            g_ban_table[i].banned_until > now) {
            count++;
        }
    }
    pthread_mutex_unlock(&g_ban_mutex);
    return count;
}

int ip_ban_get_tracked_count(void) {
    if (!g_initialized) return 0;

    pthread_mutex_lock(&g_ban_mutex);
    int count = 0;
    for (int i = 0; i < IP_BAN_TABLE_SIZE; i++) {
        if (g_ban_table[i].state == IP_BAN_ACTIVE) count++;
    }
    pthread_mutex_unlock(&g_ban_mutex);
    return count;
}

// Emit a periodic log summary of banned IPs.
// Call from the connection age checker (runs every ~30s).
// Only logs when there are active bans or when the interval has elapsed.
void ip_ban_log_stats(void) {
    if (!g_initialized) return;

    static time_t last_log = 0;
    time_t now = time(NULL);

    // Log every 5 minutes
    if (now - last_log < 300) return;
    last_log = now;

    pthread_mutex_lock(&g_ban_mutex);

    int banned_count = 0;
    int tracked_count = 0;

    for (int i = 0; i < IP_BAN_TABLE_SIZE; i++) {
        if (g_ban_table[i].state != IP_BAN_ACTIVE) continue;
        tracked_count++;
        if (g_ban_table[i].banned_until > now) {
            banned_count++;
            DEBUG_WARN("IP BAN: %s banned for %ld more seconds (ban #%d, failures: %d)",
                       g_ban_table[i].ip,
                       g_ban_table[i].banned_until - now,
                       g_ban_table[i].ban_count,
                       g_ban_table[i].failure_count);
        }
    }

    pthread_mutex_unlock(&g_ban_mutex);

    if (banned_count > 0 || tracked_count > 0) {
        DEBUG_WARN("IP BAN SUMMARY: %d IPs currently banned, %d IPs tracked",
                   banned_count, tracked_count);
    }
}