#include "test_framework.h"
#include "../src/types.h"
#include "../src/rate_limit.h"
#include <string.h>

void test_rate_limit_disabled(void) {
    TEST_SUITE_BEGIN("Rate Limit Disabled");

    rate_limit_cleanup();

    TEST_ASSERT_EQ(1, rate_limit_check("192.168.1.1"), "Rate limit check passes when disabled");
    TEST_ASSERT_EQ(1, rate_limit_check("10.0.0.1"), "Any IP passes when disabled");
    TEST_ASSERT_EQ(1, rate_limit_check(NULL), "NULL IP passes when disabled");

    TEST_SUITE_END();
}

void test_rate_limit_init(void) {
    TEST_SUITE_BEGIN("Rate Limit Init");

    rate_limit_cleanup();
    rate_limit_init(100, 60);

    TEST_ASSERT_EQ(1, rate_limit_check("192.168.1.100"), "First request passes");

    rate_limit_cleanup();

    TEST_SUITE_END();
}

void test_rate_limit_within_limit(void) {
    TEST_SUITE_BEGIN("Rate Limit Within Limit");

    rate_limit_cleanup();
    rate_limit_init(10, 60);

    const char *ip = "192.168.1.50";
    int all_passed = 1;

    for (int i = 0; i < 10; i++) {
        if (rate_limit_check(ip) != 1) {
            all_passed = 0;
            break;
        }
    }

    TEST_ASSERT_EQ(1, all_passed, "All requests within limit pass");

    rate_limit_cleanup();

    TEST_SUITE_END();
}

void test_rate_limit_exceeded(void) {
    TEST_SUITE_BEGIN("Rate Limit Exceeded");

    rate_limit_cleanup();
    rate_limit_init(5, 60);

    const char *ip = "192.168.1.200";

    for (int i = 0; i < 5; i++) {
        rate_limit_check(ip);
    }

    int exceeded = rate_limit_check(ip);
    TEST_ASSERT_EQ(0, exceeded, "Request exceeding limit is blocked");

    rate_limit_cleanup();

    TEST_SUITE_END();
}

void test_rate_limit_different_ips(void) {
    TEST_SUITE_BEGIN("Rate Limit Different IPs");

    rate_limit_cleanup();
    rate_limit_init(3, 60);

    for (int i = 0; i < 3; i++) {
        rate_limit_check("10.0.0.1");
    }
    TEST_ASSERT_EQ(0, rate_limit_check("10.0.0.1"), "First IP blocked after limit");

    TEST_ASSERT_EQ(1, rate_limit_check("10.0.0.2"), "Second IP still allowed");
    TEST_ASSERT_EQ(1, rate_limit_check("10.0.0.3"), "Third IP still allowed");

    rate_limit_cleanup();

    TEST_SUITE_END();
}

void test_rate_limit_null_ip(void) {
    TEST_SUITE_BEGIN("Rate Limit NULL IP");

    rate_limit_cleanup();
    rate_limit_init(5, 60);

    TEST_ASSERT_EQ(1, rate_limit_check(NULL), "NULL IP passes (fail-open)");

    rate_limit_cleanup();

    TEST_SUITE_END();
}

void test_rate_limit_purge_expired(void) {
    TEST_SUITE_BEGIN("Rate Limit Purge Expired");

    rate_limit_cleanup();
    rate_limit_init(100, 60);

    rate_limit_check("192.168.1.1");
    rate_limit_check("192.168.1.2");
    rate_limit_check("192.168.1.3");

    rate_limit_purge_expired();

    TEST_ASSERT_EQ(1, rate_limit_check("192.168.1.1"), "Entry still valid after purge");

    rate_limit_cleanup();

    TEST_SUITE_END();
}

void test_rate_limit_cleanup(void) {
    TEST_SUITE_BEGIN("Rate Limit Cleanup");

    rate_limit_init(10, 60);

    for (int i = 0; i < 100; i++) {
        char ip[32];
        snprintf(ip, sizeof(ip), "192.168.%d.%d", i / 256, i % 256);
        rate_limit_check(ip);
    }

    rate_limit_cleanup();

    TEST_ASSERT_EQ(1, rate_limit_check("192.168.1.1"), "After cleanup, rate limit disabled");

    TEST_SUITE_END();
}

void run_rate_limit_tests(void) {
    test_rate_limit_disabled();
    test_rate_limit_init();
    test_rate_limit_within_limit();
    test_rate_limit_exceeded();
    test_rate_limit_different_ips();
    test_rate_limit_null_ip();
    test_rate_limit_purge_expired();
    test_rate_limit_cleanup();
}