#include "a2s.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <arpa/inet.h>

/*
 * A2S_INFO — Valve Source Engine server query protocol (UDP).
 *
 * Reference: https://developer.valvesoftware.com/wiki/Server_queries
 *
 * Wire protocol (modern Source, post-Orange Box):
 *
 *   Request:
 *     FF FF FF FF 54  "Source Engine Query" 00
 *
 *   Some servers issue a challenge before responding:
 *     FF FF FF FF 41 <4-byte challenge>
 *   In that case resend the request with the challenge appended.
 *
 *   Response (type 0x49 = 'I'):
 *     FF FF FF FF 49
 *     <byte>   protocol
 *     <string> name
 *     <string> map          ← we want this
 *     <string> folder
 *     <string> game
 *     <short>  appid
 *     <byte>   players      ← we want this
 *     <byte>   max_players  ← we want this
 *     ... (remaining fields not parsed)
 */

/* ── helpers ─────────────────────────────────────────────────────────────── */

/* Read a null-terminated string from buf starting at *pos.
 * Advances *pos past the null byte.
 * Returns 0 on success, -1 if the buffer would be overrun. */
static int read_string(const uint8_t *buf, int buf_len,
                       int *pos, char *out, int out_size)
{
    int start = *pos;
    while (*pos < buf_len && buf[*pos] != '\0')
        (*pos)++;
    if (*pos >= buf_len)
        return -1;  /* no null terminator found */

    int len = *pos - start;
    if (len >= out_size)
        len = out_size - 1;
    memcpy(out, buf + start, len);
    out[len] = '\0';

    (*pos)++;  /* skip null byte */
    return 0;
}

/* ── a2s_query ───────────────────────────────────────────────────────────── */

int a2s_query(const char *host, uint16_t port, int timeout_ms, A2SInfo *out)
{
    memset(out, 0, sizeof(*out));
    out->players = -1;
    out->max_players = -1;

    /* Resolve host */
    struct addrinfo hints, *res, *rp;
    memset(&hints, 0, sizeof(hints));
    hints.ai_family   = AF_UNSPEC;
    hints.ai_socktype = SOCK_DGRAM;

    char port_str[8];
    snprintf(port_str, sizeof(port_str), "%u", port);

    if (getaddrinfo(host, port_str, &hints, &res) != 0)
        return -1;

    int fd = -1;
    for (rp = res; rp; rp = rp->ai_next) {
        fd = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
        if (fd < 0) continue;
        if (connect(fd, rp->ai_addr, rp->ai_addrlen) == 0) break;
        close(fd);
        fd = -1;
    }
    freeaddrinfo(res);
    if (fd < 0) return -1;

    /* Set receive timeout */
    struct timeval tv;
    tv.tv_sec  = timeout_ms / 1000;
    tv.tv_usec = (timeout_ms % 1000) * 1000;
    setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
    setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv));

    /* Build initial A2S_INFO request */
    static const uint8_t a2s_prefix[] = {
        0xFF, 0xFF, 0xFF, 0xFF, 0x54
    };
    static const char a2s_payload[] = "Source Engine Query";
    /* Full request: prefix + payload + null byte */
    uint8_t req[32];
    int req_len = 0;
    memcpy(req + req_len, a2s_prefix, sizeof(a2s_prefix));
    req_len += sizeof(a2s_prefix);
    memcpy(req + req_len, a2s_payload, sizeof(a2s_payload)); /* includes \0 */
    req_len += sizeof(a2s_payload);

    uint8_t resp[2048];
    int resp_len;
    int attempts = 0;

retry:
    if (attempts++ >= 2) {
        close(fd);
        return -1;
    }

    if (send(fd, req, req_len, 0) != req_len) {
        close(fd);
        return -1;
    }

    resp_len = recv(fd, resp, sizeof(resp), 0);
    if (resp_len < 5) {
        close(fd);
        return -1;
    }

    /* Check for challenge response: FF FF FF FF 41 <4 bytes> */
    if (resp_len >= 9 &&
        resp[0] == 0xFF && resp[1] == 0xFF &&
        resp[2] == 0xFF && resp[3] == 0xFF &&
        resp[4] == 0x41)
    {
        /* Append the 4-byte challenge to the request and resend */
        if (req_len + 4 > (int)sizeof(req)) {
            close(fd);
            return -1;
        }
        memcpy(req + req_len, resp + 5, 4);
        req_len += 4;
        goto retry;
    }

    /* Expect response type 0x49 ('I') */
    if (resp[0] != 0xFF || resp[1] != 0xFF ||
        resp[2] != 0xFF || resp[3] != 0xFF ||
        resp[4] != 0x49)
    {
        close(fd);
        return -1;
    }

    close(fd);

    /* Parse response body starting after the 5-byte header */
    int pos = 5;

    /* Skip protocol version (1 byte) */
    if (pos + 1 > resp_len) return -1;
    pos++;

    /* Skip name string */
    char scratch[256];
    if (read_string(resp, resp_len, &pos, scratch, sizeof(scratch)) != 0)
        return -1;

    /* Map string — this is what we want */
    if (read_string(resp, resp_len, &pos, out->map, sizeof(out->map)) != 0)
        return -1;

    /* Skip folder string */
    if (read_string(resp, resp_len, &pos, scratch, sizeof(scratch)) != 0)
        return -1;

    /* Skip game string */
    if (read_string(resp, resp_len, &pos, scratch, sizeof(scratch)) != 0)
        return -1;

    /* Skip app ID (2 bytes) */
    if (pos + 2 > resp_len) return -1;
    pos += 2;

    /* Players (1 byte) */
    if (pos + 1 > resp_len) return -1;
    out->players = resp[pos++];

    /* Max players (1 byte) */
    if (pos + 1 > resp_len) return -1;
    out->max_players = resp[pos++];

    return 0;
}