#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).
//
// State is persisted to the ip_bans SQLite table every 5 minutes
// and loaded at startup so bans survive relay restarts.
// ============================================================

#define IP_BAN_EMPTY    0
#define IP_BAN_ACTIVE   1

typedef struct {
    int state;                      // IP_BAN_EMPTY or IP_BAN_ACTIVE
    char ip[46];                    // IPv4 or IPv6 string
    
    // Auth failure tracking (existing)
    int failure_count;              // failures in current window
    time_t first_failure;           // start of current failure window
    time_t banned_until;            // 0 = not banned
    int ban_count;                  // escalation level (for exponential backoff)
    
    // NEW: Idle failure tracking (separate)
    int idle_failure_count;
    time_t idle_first_failure;
    time_t idle_banned_until;
    int idle_ban_count;
    
    // Other existing fields
    int has_authed_successfully;    // 1 if this IP has ever authenticated
    time_t last_success_at;         // timestamp of last successful auth
    int total_connections;          // lifetime connection count
    int total_failures;             // lifetime auth failure count
    int total_successes;            // lifetime successful auth count
    time_t first_seen;              // when this IP was first seen
} 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)
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;
        }
        if (strcmp(g_ban_table[idx].ip, ip) == 0) {
            return (int)idx;
        }
    }
    return -1;
}

// Get or create an entry for an IP. Returns NULL if table is full.
// Caller must hold g_ban_mutex.
static ip_ban_entry_t* get_or_create_entry(const char* ip) {
    int idx = find_slot(ip);
    if (idx < 0) {
        DEBUG_WARN("IP ban table full, cannot track %s", ip);
        return NULL;
    }
    ip_ban_entry_t* entry = &g_ban_table[idx];
    if (entry->state == IP_BAN_EMPTY) {
        entry->state = IP_BAN_ACTIVE;
        strncpy(entry->ip, ip, sizeof(entry->ip) - 1);
        entry->ip[sizeof(entry->ip) - 1] = '\0';
        entry->first_seen = time(NULL);
    }
    return entry;
}

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);
}

void ip_ban_load_from_db(sqlite3* db) {
    if (!db || !g_initialized) return;

    // Create table if it doesn't exist (handles existing databases)
    const char* create_sql =
        "CREATE TABLE IF NOT EXISTS ip_bans ("
        "  ip TEXT PRIMARY KEY,"
        "  failure_count INTEGER NOT NULL DEFAULT 0,"
        "  ban_count INTEGER NOT NULL DEFAULT 0,"
        "  banned_until INTEGER NOT NULL DEFAULT 0,"
        "  first_failure INTEGER NOT NULL DEFAULT 0,"
        "  has_authed_successfully INTEGER NOT NULL DEFAULT 0,"
        "  last_success_at INTEGER NOT NULL DEFAULT 0,"
        "  total_connections INTEGER NOT NULL DEFAULT 0,"
        "  total_failures INTEGER NOT NULL DEFAULT 0,"
        "  total_successes INTEGER NOT NULL DEFAULT 0,"
        "  first_seen INTEGER NOT NULL DEFAULT 0,"
        "  updated_at INTEGER NOT NULL DEFAULT (strftime('%s', 'now'))"
        ");";
    char* err = NULL;
    if (sqlite3_exec(db, create_sql, NULL, NULL, &err) != SQLITE_OK) {
        DEBUG_ERROR("Failed to create ip_bans table: %s", err ? err : "unknown");
        if (err) sqlite3_free(err);
        return;
    }

    // Migration: Add idle_* columns if they don't exist (ignore errors if already exists)
    sqlite3_exec(db, "ALTER TABLE ip_bans ADD COLUMN idle_failure_count INTEGER NOT NULL DEFAULT 0", NULL, NULL, NULL);
    sqlite3_exec(db, "ALTER TABLE ip_bans ADD COLUMN idle_ban_count INTEGER NOT NULL DEFAULT 0", NULL, NULL, NULL);
    sqlite3_exec(db, "ALTER TABLE ip_bans ADD COLUMN idle_banned_until INTEGER NOT NULL DEFAULT 0", NULL, NULL, NULL);
    sqlite3_exec(db, "ALTER TABLE ip_bans ADD COLUMN idle_first_failure INTEGER NOT NULL DEFAULT 0", NULL, NULL, NULL);

    const char* sql =
        "SELECT ip, failure_count, ban_count, banned_until, first_failure,"
        "       has_authed_successfully, last_success_at, total_connections,"
        "       total_failures, total_successes, first_seen,"
        "       idle_failure_count, idle_ban_count, idle_banned_until, idle_first_failure"
        " FROM ip_bans";

    sqlite3_stmt* stmt;
    if (sqlite3_prepare_v2(db, sql, -1, &stmt, NULL) != SQLITE_OK) {
        DEBUG_ERROR("Failed to prepare ip_bans load query");
        return;
    }

    int loaded = 0;
    pthread_mutex_lock(&g_ban_mutex);
    while (sqlite3_step(stmt) == SQLITE_ROW) {
        const char* ip = (const char*)sqlite3_column_text(stmt, 0);
        if (!ip) continue;

        int idx = find_slot(ip);
        if (idx < 0) continue;

        ip_ban_entry_t* entry = &g_ban_table[idx];
        entry->state = IP_BAN_ACTIVE;
        strncpy(entry->ip, ip, sizeof(entry->ip) - 1);
        entry->ip[sizeof(entry->ip) - 1] = '\0';
        entry->failure_count          = sqlite3_column_int(stmt, 1);
        entry->ban_count              = sqlite3_column_int(stmt, 2);
        entry->banned_until           = (time_t)sqlite3_column_int64(stmt, 3);
        entry->first_failure          = (time_t)sqlite3_column_int64(stmt, 4);
        entry->has_authed_successfully = sqlite3_column_int(stmt, 5);
        entry->last_success_at        = (time_t)sqlite3_column_int64(stmt, 6);
        entry->total_connections      = sqlite3_column_int(stmt, 7);
        entry->total_failures         = sqlite3_column_int(stmt, 8);
        entry->total_successes        = sqlite3_column_int(stmt, 9);
        entry->first_seen             = (time_t)sqlite3_column_int64(stmt, 10);
        // Load idle tracking fields (default to 0 if columns don't exist yet)
        entry->idle_failure_count     = sqlite3_column_int(stmt, 11);
        entry->idle_ban_count         = sqlite3_column_int(stmt, 12);
        entry->idle_banned_until      = (time_t)sqlite3_column_int64(stmt, 13);
        entry->idle_first_failure     = (time_t)sqlite3_column_int64(stmt, 14);
        loaded++;
    }
    pthread_mutex_unlock(&g_ban_mutex);
    sqlite3_finalize(stmt);

    // Count how many are still actively banned
    time_t now = time(NULL);
    int still_banned = 0;
    pthread_mutex_lock(&g_ban_mutex);
    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) {
            still_banned++;
        }
    }
    pthread_mutex_unlock(&g_ban_mutex);

    DEBUG_WARN("IP ban table loaded: %d IPs restored (%d still banned)", loaded, still_banned);
}

void ip_ban_save_to_db(sqlite3* db) {
    if (!db || !g_initialized) return;

    const char* upsert_sql =
        "INSERT OR REPLACE INTO ip_bans"
        " (ip, failure_count, ban_count, banned_until, first_failure,"
        "  has_authed_successfully, last_success_at, total_connections,"
        "  total_failures, total_successes, first_seen, updated_at,"
        "  idle_failure_count, idle_ban_count, idle_banned_until, idle_first_failure)"
        " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, strftime('%s', 'now'), ?, ?, ?, ?)";

    sqlite3_stmt* stmt;
    if (sqlite3_prepare_v2(db, upsert_sql, -1, &stmt, NULL) != SQLITE_OK) {
        DEBUG_ERROR("Failed to prepare ip_bans save query");
        return;
    }

    sqlite3_exec(db, "BEGIN TRANSACTION", NULL, NULL, NULL);

    int saved = 0;
    pthread_mutex_lock(&g_ban_mutex);
    for (int i = 0; i < IP_BAN_TABLE_SIZE; i++) {
        ip_ban_entry_t* entry = &g_ban_table[i];
        if (entry->state != IP_BAN_ACTIVE) continue;

        sqlite3_reset(stmt);
        sqlite3_bind_text(stmt, 1, entry->ip, -1, SQLITE_STATIC);
        sqlite3_bind_int(stmt, 2, entry->failure_count);
        sqlite3_bind_int(stmt, 3, entry->ban_count);
        sqlite3_bind_int64(stmt, 4, (sqlite3_int64)entry->banned_until);
        sqlite3_bind_int64(stmt, 5, (sqlite3_int64)entry->first_failure);
        sqlite3_bind_int(stmt, 6, entry->has_authed_successfully);
        sqlite3_bind_int64(stmt, 7, (sqlite3_int64)entry->last_success_at);
        sqlite3_bind_int(stmt, 8, entry->total_connections);
        sqlite3_bind_int(stmt, 9, entry->total_failures);
        sqlite3_bind_int(stmt, 10, entry->total_successes);
        sqlite3_bind_int64(stmt, 11, (sqlite3_int64)entry->first_seen);
        // Idle tracking fields
        sqlite3_bind_int(stmt, 12, entry->idle_failure_count);
        sqlite3_bind_int(stmt, 13, entry->idle_ban_count);
        sqlite3_bind_int64(stmt, 14, (sqlite3_int64)entry->idle_banned_until);
        sqlite3_bind_int64(stmt, 15, (sqlite3_int64)entry->idle_first_failure);

        if (sqlite3_step(stmt) != SQLITE_DONE) {
            DEBUG_WARN("Failed to save ip_ban entry for %s", entry->ip);
        } else {
            saved++;
        }
    }
    pthread_mutex_unlock(&g_ban_mutex);

    sqlite3_finalize(stmt);
    sqlite3_exec(db, "COMMIT", NULL, NULL, NULL);

    DEBUG_TRACE("IP ban table saved: %d entries written to DB", saved);
}

// Check if an IP is in the idle_ban_whitelist config (comma-separated list)
static int ip_is_whitelisted(const char* ip) {
    const char* whitelist = get_config_value("idle_ban_whitelist");
    if (!whitelist || whitelist[0] == '\0') return 0;

    // Make a mutable copy to tokenize
    char buf[1024];
    strncpy(buf, whitelist, sizeof(buf) - 1);
    buf[sizeof(buf) - 1] = '\0';

    char* token = strtok(buf, ",");
    while (token) {
        // Trim leading/trailing spaces
        while (*token == ' ') token++;
        char* end = token + strlen(token) - 1;
        while (end > token && *end == ' ') { *end = '\0'; end--; }
        if (strcmp(token, ip) == 0) return 1;
        token = strtok(NULL, ",");
    }
    return 0;
}

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

    // Whitelisted IPs are never banned
    if (ip_is_whitelisted(ip)) 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;
    }

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

    // Check auth ban (if enabled)
    if (get_config_bool("auth_fail_ban_enabled", 1) &&
        entry->banned_until > 0 && now < entry->banned_until) {
        banned = 1;
    }

    // Check idle ban (if enabled)
    if (get_config_bool("idle_ban_enabled", 1) &&
        entry->idle_banned_until > 0 && now < entry->idle_banned_until) {
        banned = 1;
    }

    // Clear expired bans
    if (!banned) {
        if (entry->banned_until > 0 && now >= entry->banned_until) {
            entry->banned_until = 0;
            entry->failure_count = 0;
            entry->first_failure = 0;
        }
        if (entry->idle_banned_until > 0 && now >= entry->idle_banned_until) {
            entry->idle_banned_until = 0;
            entry->idle_failure_count = 0;
            entry->idle_first_failure = 0;
        }
    }

    pthread_mutex_unlock(&g_ban_mutex);
    return banned;
}

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

    pthread_mutex_lock(&g_ban_mutex);
    ip_ban_entry_t* entry = get_or_create_entry(ip);
    if (entry) {
        entry->total_connections++;
    }
    pthread_mutex_unlock(&g_ban_mutex);
}

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);
    ip_ban_entry_t* entry = get_or_create_entry(ip);
    if (!entry) {
        pthread_mutex_unlock(&g_ban_mutex);
        return;
    }

    time_t now = time(NULL);

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

    entry->failure_count++;
    entry->total_failures++;

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

    if (entry->failure_count >= threshold) {
        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;
        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);
}

// Record an idle/early-disconnect failure for an IP
void ip_ban_record_idle_failure(const char* ip) {
    if (!ip || !g_initialized) return;
    if (!get_config_bool("idle_ban_enabled", 1)) return;
    if (ip_is_whitelisted(ip)) return;  // Never record idle failures for whitelisted IPs

    int threshold = get_config_int("idle_ban_threshold", 1);
    int window_sec = get_config_int("idle_ban_window_sec", 30);
    int ban_duration = get_config_int("idle_ban_duration_sec", 300);

    pthread_mutex_lock(&g_ban_mutex);
    ip_ban_entry_t* entry = get_or_create_entry(ip);
    if (!entry) {
        pthread_mutex_unlock(&g_ban_mutex);
        return;
    }

    time_t now = time(NULL);

    // Reset window if expired
    if (entry->idle_first_failure > 0 && (now - entry->idle_first_failure) > window_sec) {
        entry->idle_failure_count = 0;
        entry->idle_first_failure = now;
    }
    if (entry->idle_first_failure == 0) {
        entry->idle_first_failure = now;
    }

    entry->idle_failure_count++;
    entry->total_failures++;

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

    if (entry->idle_failure_count >= threshold) {
        int duration = ban_duration;
        for (int i = 0; i < entry->idle_ban_count && duration < 86400; i++) {
            duration *= 2;
        }
        if (duration > 86400) duration = 86400;

        entry->idle_banned_until = now + duration;
        entry->idle_ban_count++;
        entry->idle_failure_count = 0;
        entry->idle_first_failure = 0;

        DEBUG_WARN("IP %s banned for %d seconds (idle ban #%d) after %d idle failures",
                   ip, duration, entry->idle_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);
    ip_ban_entry_t* entry = get_or_create_entry(ip);
    if (entry) {
        entry->failure_count = 0;
        entry->first_failure = 0;
        entry->has_authed_successfully = 1;
        entry->last_success_at = time(NULL);
        entry->total_successes++;
        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 idle_window_sec = get_config_int("idle_ban_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];

        // Check auth failure window expiration
        int auth_ban_expired = (entry->banned_until == 0 || now >= entry->banned_until);
        int auth_window_expired = (entry->first_failure == 0 || (now - entry->first_failure) > window_sec * 10);

        // Check idle failure window expiration
        int idle_ban_expired = (entry->idle_banned_until == 0 || now >= entry->idle_banned_until);
        int idle_window_expired = (entry->idle_first_failure == 0 || (now - entry->idle_first_failure) > idle_window_sec * 10);

        if (auth_ban_expired && auth_window_expired && entry->failure_count == 0 &&
            idle_ban_expired && idle_window_expired && entry->idle_failure_count == 0) {
            int retain_sec = 86400; // 24 hours
            int last_auth_ban_expired_long_ago = (entry->banned_until == 0 ||
                                                   (now - entry->banned_until) > retain_sec);
            int last_idle_ban_expired_long_ago = (entry->idle_banned_until == 0 ||
                                                   (now - entry->idle_banned_until) > retain_sec);

            if (last_auth_ban_expired_long_ago && last_idle_ban_expired_long_ago &&
                !entry->has_authed_successfully &&
                entry->ban_count == 0 && entry->idle_ban_count == 0 &&
                entry->total_connections <= 1) {
                // Fully clean — never banned, never authenticated, only seen once
                memset(entry, 0, sizeof(ip_ban_entry_t));
                cleaned++;
            } else {
                // Keep entry permanently — preserve ban_count for escalation.
                // An IP that has been banned before will always get at least a 24-hour ban
                // if it fails auth again, regardless of how long it has been away.
                entry->failure_count = 0;
                entry->first_failure = 0;
                entry->idle_failure_count = 0;
                entry->idle_first_failure = 0;
                if (last_auth_ban_expired_long_ago) {
                    entry->banned_until = 0;
                    // ban_count intentionally NOT reset — permanent escalation
                }
                if (last_idle_ban_expired_long_ago) {
                    entry->idle_banned_until = 0;
                    // idle_ban_count intentionally NOT reset — permanent escalation
                }
            }
        }
    }

    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) continue;
        // Count if either auth banned or idle banned
        if (g_ban_table[i].banned_until > now ||
            g_ban_table[i].idle_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;
}

void ip_ban_log_stats(sqlite3* db) {
    if (!g_initialized) return;

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

    if (now - last_log < 300) return;
    last_log = now;

    // Save to DB every 5 minutes
    if (db) {
        ip_ban_save_to_db(db);
    }

    pthread_mutex_lock(&g_ban_mutex);
    int auth_banned_count = 0;
    int idle_banned_count = 0;
    int tracked_count = 0;
    int trusted_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) auth_banned_count++;
        if (g_ban_table[i].idle_banned_until > now) idle_banned_count++;
        if (g_ban_table[i].has_authed_successfully) trusted_count++;
    }
    pthread_mutex_unlock(&g_ban_mutex);

    int total_banned = auth_banned_count + idle_banned_count;
    if (total_banned > 0 || tracked_count > 0) {
        DEBUG_WARN("IP BAN SUMMARY: %d auth-banned, %d idle-banned, %d tracked, %d trusted (ever authed)",
                   auth_banned_count, idle_banned_count, tracked_count, trusted_count);
    }
}