#include "health_check.h"
#include "logging.h"
#include "config.h"
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <fcntl.h>
#include <errno.h>
#include <poll.h>
#include <pthread.h>

typedef struct {
    char hostname[256];
    char upstream_host[256];
    int upstream_port;
    int healthy;
    int consecutive_failures;
    time_t last_check;
} upstream_health_t;

static upstream_health_t *health_states = NULL;
static int health_state_count = 0;
static pthread_mutex_t health_mutex = PTHREAD_MUTEX_INITIALIZER;
static int g_health_check_enabled = 0;

void health_check_init(void) {
    pthread_mutex_lock(&health_mutex);

    if (health_states) {
        free(health_states);
    }

    health_state_count = config.route_count;
    if (health_state_count <= 0) {
        health_states = NULL;
        pthread_mutex_unlock(&health_mutex);
        return;
    }

    health_states = calloc(health_state_count, sizeof(upstream_health_t));
    if (!health_states) {
        health_state_count = 0;
        pthread_mutex_unlock(&health_mutex);
        return;
    }

    for (int i = 0; i < health_state_count; i++) {
        strncpy(health_states[i].hostname, config.routes[i].hostname, sizeof(health_states[i].hostname) - 1);
        strncpy(health_states[i].upstream_host, config.routes[i].upstream_host, sizeof(health_states[i].upstream_host) - 1);
        health_states[i].upstream_port = config.routes[i].upstream_port;
        health_states[i].healthy = 1;
        health_states[i].consecutive_failures = 0;
        health_states[i].last_check = 0;
    }

    g_health_check_enabled = 1;
    log_info("Health check initialized for %d upstreams", health_state_count);

    pthread_mutex_unlock(&health_mutex);
}

void health_check_cleanup(void) {
    pthread_mutex_lock(&health_mutex);
    if (health_states) {
        free(health_states);
        health_states = NULL;
    }
    health_state_count = 0;
    g_health_check_enabled = 0;
    pthread_mutex_unlock(&health_mutex);
}

static int check_tcp_connection(const char *host, int port, int timeout_ms) {
    struct addrinfo hints, *result;
    memset(&hints, 0, sizeof(hints));
    hints.ai_family = AF_INET;
    hints.ai_socktype = SOCK_STREAM;

    if (getaddrinfo(host, NULL, &hints, &result) != 0) {
        return 0;
    }

    int fd = socket(AF_INET, SOCK_STREAM, 0);
    if (fd < 0) {
        freeaddrinfo(result);
        return 0;
    }

    int flags = fcntl(fd, F_GETFL, 0);
    fcntl(fd, F_SETFL, flags | O_NONBLOCK);

    struct sockaddr_in *addr = (struct sockaddr_in *)result->ai_addr;
    addr->sin_port = htons(port);

    int connect_result = connect(fd, (struct sockaddr *)addr, sizeof(struct sockaddr_in));
    freeaddrinfo(result);

    if (connect_result == 0) {
        close(fd);
        return 1;
    }

    if (errno != EINPROGRESS) {
        close(fd);
        return 0;
    }

    struct pollfd pfd;
    pfd.fd = fd;
    pfd.events = POLLOUT;

    int poll_result = poll(&pfd, 1, timeout_ms);
    if (poll_result <= 0) {
        close(fd);
        return 0;
    }

    int error = 0;
    socklen_t len = sizeof(error);
    getsockopt(fd, SOL_SOCKET, SO_ERROR, &error, &len);
    close(fd);

    return error == 0;
}

void health_check_run(void) {
    if (!g_health_check_enabled) return;

    pthread_mutex_lock(&health_mutex);

    time_t now = time(NULL);

    for (int i = 0; i < health_state_count; i++) {
        if (now - health_states[i].last_check < HEALTH_CHECK_INTERVAL_SECONDS) {
            continue;
        }

        health_states[i].last_check = now;

        int is_healthy = check_tcp_connection(
            health_states[i].upstream_host,
            health_states[i].upstream_port,
            HEALTH_CHECK_TIMEOUT_MS
        );

        if (is_healthy) {
            if (!health_states[i].healthy) {
                log_info("Upstream %s:%d is now healthy",
                         health_states[i].upstream_host,
                         health_states[i].upstream_port);
            }
            health_states[i].healthy = 1;
            health_states[i].consecutive_failures = 0;
        } else {
            health_states[i].consecutive_failures++;
            if (health_states[i].consecutive_failures >= 3 && health_states[i].healthy) {
                log_info("Upstream %s:%d is now unhealthy (failures: %d)",
                         health_states[i].upstream_host,
                         health_states[i].upstream_port,
                         health_states[i].consecutive_failures);
                health_states[i].healthy = 0;
            }
        }
    }

    pthread_mutex_unlock(&health_mutex);
}

int health_check_is_healthy(const char *hostname) {
    if (!g_health_check_enabled || !hostname) return 1;

    pthread_mutex_lock(&health_mutex);

    for (int i = 0; i < health_state_count; i++) {
        if (strcasecmp(health_states[i].hostname, hostname) == 0) {
            int result = health_states[i].healthy;
            pthread_mutex_unlock(&health_mutex);
            return result;
        }
    }

    pthread_mutex_unlock(&health_mutex);
    return 1;
}