#include "client.h"
#include "proto.h"
#include "unit.h"
#include "log.h"
#include "umbrella.h"
#include "console/rcon.h"
#include "console/a2s.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <sys/epoll.h>
#include <sys/stat.h>
#include <systemd/sd-bus.h>

/* ── Socket setup ────────────────────────────────────────────────────────── */

int client_listen(void) {
    /* Remove stale socket if present */
    unlink(UMBRELLA_SOCK_PATH);

    int fd = socket(AF_UNIX, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
    if (fd < 0) {
        log_error("socket: %s", strerror(errno));
        return -1;
    }

    struct sockaddr_un addr;
    memset(&addr, 0, sizeof(addr));
    addr.sun_family = AF_UNIX;
    strncpy(addr.sun_path, UMBRELLA_SOCK_PATH, sizeof(addr.sun_path) - 1);

    if (bind(fd, (struct sockaddr *)&addr, sizeof(addr)) != 0) {
        log_error("bind %s: %s", UMBRELLA_SOCK_PATH, strerror(errno));
        close(fd);
        return -1;
    }

    /* Owner read/write, group read/write, others nothing */
    chmod(UMBRELLA_SOCK_PATH, 0660);

    if (listen(fd, 16) != 0) {
        log_error("listen: %s", strerror(errno));
        close(fd);
        return -1;
    }

    log_info("Listening on %s", UMBRELLA_SOCK_PATH);
    return fd;
}

/* ── Client lifecycle ────────────────────────────────────────────────────── */

int client_accept(int listen_fd) {
    if (g.client_count >= MAX_CLIENTS) {
        log_warn("Max clients reached, rejecting connection");
        /* Accept and immediately close */
        int tmp = accept(listen_fd, NULL, NULL);
        if (tmp >= 0) close(tmp);
        return -1;
    }

    int fd = accept4(listen_fd, NULL, NULL, SOCK_NONBLOCK | SOCK_CLOEXEC);
    if (fd < 0) {
        if (errno != EAGAIN && errno != EWOULDBLOCK)
            log_error("accept: %s", strerror(errno));
        return -1;
    }

    /* Get peer credentials for access control */
    struct ucred cred;
    socklen_t cred_len = sizeof(cred);
    int uid = -1;
    if (getsockopt(fd, SOL_SOCKET, SO_PEERCRED, &cred, &cred_len) == 0)
        uid = (int)cred.uid;

    /* Find a free client slot */
    int slot = -1;
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (g.clients[i].fd == 0) { slot = i; break; }
    }
    if (slot < 0) {
        close(fd);
        return -1;
    }

    g.clients[slot].fd = fd;
    g.clients[slot].uid = uid;
    g.clients[slot].attached[0] = '\0';
    g.client_count++;

    /* Register with epoll */
    struct epoll_event ev;
    ev.events = EPOLLIN | EPOLLRDHUP;
    ev.data.fd = fd;
    epoll_ctl(g.epoll_fd, EPOLL_CTL_ADD, fd, &ev);

    log_debug("Client connected (fd=%d uid=%d)", fd, uid);
    return 0;
}

Client *client_find(int fd) {
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (g.clients[i].fd == fd)
            return &g.clients[i];
    }
    return NULL;
}

void client_remove(int fd) {
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (g.clients[i].fd == fd) {
            log_debug("Client disconnected (fd=%d)", fd);
            epoll_ctl(g.epoll_fd, EPOLL_CTL_DEL, fd, NULL);
            close(fd);
            memset(&g.clients[i], 0, sizeof(Client));
            g.client_count--;
            return;
        }
    }
}

void client_broadcast_output(const char *unit_name,
                             const char *data, int history) {
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (g.clients[i].fd == 0) continue;
        if (strcmp(g.clients[i].attached, unit_name) == 0)
            proto_send_output(g.clients[i].fd, unit_name, data, history);
    }
}

/* ── Command handlers ────────────────────────────────────────────────────── */

/*
 * systemd_active_state: query the ActiveState of a systemd unit via D-Bus.
 *
 * Returns a pointer to a static string: "active", "activating",
 * "deactivating", "inactive", "failed", or "unknown" on error.
 *
 * Only called when u->service[0] is set. Works for any unit type
 * (RCON or STDIN) that declares a systemd service.
 */
static const char *systemd_active_state(const char *service_name)
{
    sd_bus *bus = NULL;
    sd_bus_message *reply = NULL;
    sd_bus_error error = SD_BUS_ERROR_NULL;
    const char *obj_path = NULL;
    char *state_str = NULL;
    const char *result = "unknown";

    if (sd_bus_open_system(&bus) < 0)
        goto out;

    /* GetUnit(service_name) → object path */
    if (sd_bus_call_method(
            bus,
            "org.freedesktop.systemd1",
            "/org/freedesktop/systemd1",
            "org.freedesktop.systemd1.Manager",
            "GetUnit",
            &error, &reply,
            "s", service_name) < 0)
        goto out;

    if (sd_bus_message_read(reply, "o", &obj_path) < 0 || !obj_path)
        goto out;

    /* Query ActiveState property on the unit object */
    if (sd_bus_get_property_string(
            bus,
            "org.freedesktop.systemd1",
            obj_path,
            "org.freedesktop.systemd1.Unit",
            "ActiveState",
            &error, &state_str) < 0)
        goto out;

    /* Map to one of our known strings */
    if      (strcmp(state_str, "active")       == 0) result = "active";
    else if (strcmp(state_str, "activating")   == 0) result = "activating";
    else if (strcmp(state_str, "deactivating") == 0) result = "deactivating";
    else if (strcmp(state_str, "inactive")     == 0) result = "inactive";
    else if (strcmp(state_str, "failed")       == 0) result = "failed";

out:
    free(state_str);
    sd_bus_message_unref(reply);
    sd_bus_error_free(&error);
    sd_bus_unref(bus);
    return result;
}

/*
 * ProbeResult: returned by probe_unit_state().
 * Carries richer information than a plain state string.
 */
typedef struct {
    const char *state;   /* pointer to static string */
    int  players;        /* -1 = not available (non-A2S unit) */
    int  max_players;    /* -1 = not available */
    char map[64];        /* empty = not available */
} ProbeResult;

/*
 * probe_unit_state: determine current state of a unit using all available
 * probing methods, honouring per-unit configuration.
 *
 * Decision tree:
 *
 *   STDIN unit:
 *     → {state = internal process state}   (tracked directly by daemon)
 *
 *   Any unit with service[0] set → sd-bus D-Bus:
 *     "activating"   → {state="starting"}
 *     "deactivating" → {state="stopping"}
 *     "inactive"     → {state="stopped"}
 *     "failed"       → {state="crashed"}
 *     "active"       → continue to application-level probe
 *
 *   RCON unit, health.type == HEALTH_A2S → A2S UDP query:
 *     success, players == 0  → {state="hibernating", players, max_players, map}
 *     success, players  > 0  → {state="running",     players, max_players, map}
 *     failure → RCON echo probe:
 *       success → {state="changing_map"}
 *       failure → {state="unreachable"}
 *
 *   RCON unit, health.type == HEALTH_NONE → RCON echo probe:
 *     success → {state="running"}
 *     failure → {state="unreachable"}
 */
static ProbeResult probe_unit_state(Unit *u)
{
    ProbeResult r = { .state = "unknown", .players = -1, .max_players = -1 };
    r.map[0] = '\0';

    /* STDIN units: state is maintained directly by the process layer */
    if (u->console.type != CONSOLE_RCON) {
        r.state = unit_state_str(u->state);
        return r;
    }

    /* For any unit (RCON or STDIN) with a systemd service, check D-Bus first */
    if (u->service[0]) {
        const char *sd = systemd_active_state(u->service);
        if      (strcmp(sd, "activating")   == 0) { r.state = "starting";  return r; }
        else if (strcmp(sd, "deactivating") == 0) { r.state = "stopping";  return r; }
        else if (strcmp(sd, "inactive")     == 0) { r.state = "stopped";   return r; }
        else if (strcmp(sd, "failed")       == 0) { r.state = "crashed";   return r; }
        /* "active" or "unknown" → fall through to application probe */
    }

    /* A2S health probe (Valve games: TF2, GMod, …) */
    if (u->health.type == HEALTH_A2S) {
        const char *h_host = u->health.host[0] ? u->health.host : u->console.host;
        uint16_t    h_port = u->health.port     ? u->health.port : u->console.port;
        int         h_tmo  = u->health.timeout_ms > 0 ? u->health.timeout_ms : 5000;

        A2SInfo info;
        if (a2s_query(h_host, h_port, h_tmo, &info) == 0) {
            r.players     = info.players;
            r.max_players = info.max_players;
            strncpy(r.map, info.map, sizeof(r.map) - 1);
            r.state = (r.players == 0) ? "hibernating" : "running";
            return r;
        }

        /* A2S failed — try RCON to distinguish changing_map from unreachable */
        char response[64] = {0};
        int ok = rcon_exec(u->console.host, u->console.port,
                           u->console.password, "echo umbrella_probe",
                           response, sizeof(response));
        r.state = (ok == 0) ? "changing_map" : "unreachable";
        return r;
    }

    /* Default: plain RCON echo probe */
    char response[64] = {0};
    int ok = rcon_exec(u->console.host, u->console.port,
                       u->console.password, "echo umbrella_probe",
                       response, sizeof(response));
    r.state = (ok == 0) ? "running" : "unreachable";
    return r;
}

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

/* Append probe result's optional fields to an in-progress JSON buffer.
 * Call after the last fixed field; this adds players/max_players/map. */
static int append_probe_fields(char *buf, int pos, int size,
                               const ProbeResult *r)
{
    pos += snprintf(buf + pos, size - pos,
                    ",\"players\":%d,\"max_players\":%d,\"map\":\"%s\"",
                    r->players, r->max_players, r->map);
    return pos;
}

/* ── cmd_list ─────────────────────────────────────────────────────────────── */

static void cmd_list(Client *c) {
    char buf[PROTO_MAX_MSG];
    int pos = 0;

    pos += snprintf(buf + pos, sizeof(buf) - pos,
                    "{\"type\":\"list\",\"units\":[");

    for (int i = 0; i < g.unit_count; i++) {
        Unit *u = &g.units[i];
        ProbeResult r = probe_unit_state(u);

        if (i > 0)
            pos += snprintf(buf + pos, sizeof(buf) - pos, ",");

        pos += snprintf(buf + pos, sizeof(buf) - pos,
                        "{\"name\":\"%s\",\"display\":\"%s\",\"state\":\"%s\"",
                        u->name, u->display, r.state);
        pos = append_probe_fields(buf, pos, sizeof(buf), &r);
        pos += snprintf(buf + pos, sizeof(buf) - pos, "}");
    }

    snprintf(buf + pos, sizeof(buf) - pos, "]}");
    proto_send(c->fd, buf);
}

/* ── cmd_status ──────────────────────────────────────────────────────────── */

static void cmd_status(Client *c, const char *unit_name) {
    Unit *u = unit_find(unit_name);
    if (!u) { proto_send_error(c->fd, "unit not found"); return; }

    ProbeResult r = probe_unit_state(u);

    char buf[512];
    int pos = 0;
    pos += snprintf(buf + pos, sizeof(buf) - pos,
                    "{\"type\":\"status\",\"unit\":\"%s\",\"display\":\"%s\","
                    "\"state\":\"%s\",\"pid\":%d",
                    u->name, u->display, r.state, (int)u->pid);
    pos = append_probe_fields(buf, pos, sizeof(buf), &r);
    snprintf(buf + pos, sizeof(buf) - pos, "}");
    proto_send(c->fd, buf);
}

/* ── cmd_tail ────────────────────────────────────────────────────────────── */

static void cmd_tail(Client *c, const char *unit_name) {
    Unit *u = unit_find(unit_name);
    if (!u) { proto_send_error(c->fd, "unit not found"); return; }

    /* Concatenate all ring buffer lines into a single payload */
    char data[PROTO_MAX_MSG - 256];
    int dpos = 0;

    RingBuffer *rb = u->output;
    if (rb && rb->count > 0) {
        int start = (rb->head - rb->count + RING_BUF_LINES) % RING_BUF_LINES;
        for (int i = 0; i < rb->count && dpos < (int)sizeof(data) - 2; i++) {
            int idx = (start + i) % RING_BUF_LINES;
            const char *line = rb->lines[idx];
            int len = strlen(line);
            if (dpos + len + 1 >= (int)sizeof(data))
                break;
            memcpy(data + dpos, line, len);
            dpos += len;
            /* Ensure each line ends with a newline */
            if (dpos > 0 && data[dpos - 1] != '\n')
                data[dpos++] = '\n';
        }
    }
    data[dpos] = '\0';

    /* Escape the data for JSON */
    char escaped[PROTO_MAX_MSG];
    int j = 0;
    for (int i = 0; data[i] && j < (int)sizeof(escaped) - 2; i++) {
        switch (data[i]) {
            case '"':  escaped[j++] = '\\'; escaped[j++] = '"';  break;
            case '\\': escaped[j++] = '\\'; escaped[j++] = '\\'; break;
            case '\n': escaped[j++] = '\\'; escaped[j++] = 'n';  break;
            case '\r': escaped[j++] = '\\'; escaped[j++] = 'r';  break;
            default:   escaped[j++] = data[i]; break;
        }
    }
    escaped[j] = '\0';

    /* Send response envelope like other request/response commands */
    char buf[PROTO_MAX_MSG];
    snprintf(buf, sizeof(buf),
             "{\"type\":\"ok\",\"data\":\"%s\"}", escaped);
    proto_send(c->fd, buf);
}

/* ── cmd_broadcast ───────────────────────────────────────────────────────── */

static void cmd_broadcast(Client *c, const char *message) {
    int sent = 0, failed = 0;

    for (int i = 0; i < g.unit_count; i++) {
        Unit *u = &g.units[i];

        /* Skip units without a broadcast command configured */
        if (!u->broadcast_cmd[0])
            continue;

        /* Only broadcast to units the daemon considers running */
        if (u->state != STATE_RUNNING)
            continue;

        /* Substitute {msg} with the message in the broadcast command */
        char cmd[256];
        const char *tmpl = u->broadcast_cmd;
        const char *placeholder = strstr(tmpl, "{msg}");
        if (placeholder) {
            int prefix_len = (int)(placeholder - tmpl);
            snprintf(cmd, sizeof(cmd), "%.*s%s%s",
                     prefix_len, tmpl,
                     message,
                     placeholder + 5 /* strlen("{msg}") */);
        } else {
            /* No placeholder — append message after a space */
            snprintf(cmd, sizeof(cmd), "%s %s", tmpl, message);
        }

        int ok = 0;
        if (u->console.type == CONSOLE_RCON) {
            char response[256] = {0};
            ok = (rcon_exec(u->console.host, u->console.port,
                            u->console.password, cmd,
                            response, sizeof(response)) == 0);
        } else if (u->console.type == CONSOLE_STDIN) {
            if (u->stdin_fd >= 0) {
                /* Append newline so the server processes the command */
                size_t clen = strlen(cmd);
                if (clen < sizeof(cmd) - 1) {
                    cmd[clen]     = '\n';
                    cmd[clen + 1] = '\0';
                    ok = (write(u->stdin_fd, cmd, clen + 1) > 0);
                }
            }
        }

        if (ok) {
            log_info("broadcast: sent to %s: %s", u->name, cmd);
            sent++;
        } else {
            log_warn("broadcast: failed for %s", u->name);
            failed++;
        }
    }

    char buf[128];
    snprintf(buf, sizeof(buf),
             "{\"type\":\"ok\",\"sent\":%d,\"failed\":%d}", sent, failed);
    proto_send(c->fd, buf);
}

/* ── cmd_attach ──────────────────────────────────────────────────────────── */

static void cmd_attach(Client *c, const char *unit_name) {
    Unit *u = unit_find(unit_name);
    if (!u) { proto_send_error(c->fd, "unit not found"); return; }

    strncpy(c->attached, unit_name, MAX_NAME - 1);
    log_info("Client fd=%d attached to %s", c->fd, unit_name);

    /* Send ring buffer history first */
    RingBuffer *rb = u->output;
    if (rb && rb->count > 0) {
        int start = (rb->head - rb->count + RING_BUF_LINES) % RING_BUF_LINES;
        for (int i = 0; i < rb->count; i++) {
            int idx = (start + i) % RING_BUF_LINES;
            proto_send_output(c->fd, unit_name, rb->lines[idx], 1);
        }
    }

    proto_send_ok(c->fd);
}

static void cmd_detach(Client *c) {
    if (c->attached[0])
        log_info("Client fd=%d detached from %s", c->fd, c->attached);
    c->attached[0] = '\0';
    proto_send_ok(c->fd);
}

static void cmd_input(Client *c, const char *unit_name, const char *data) {
    Unit *u = unit_find(unit_name);
    if (!u) { proto_send_error(c->fd, "unit not found"); return; }

    if (u->console.type == CONSOLE_RCON) {
        /* Send via RCON — open a connection, send, close */
        char response[4096] = {0};
        int r = rcon_exec(u->console.host, u->console.port,
                          u->console.password, data,
                          response, sizeof(response));
        if (r != 0) {
            proto_send_error(c->fd, "RCON command failed");
            return;
        }
        /* Echo the response back as output */
        if (response[0])
            proto_send_output(c->fd, unit_name, response, 0);
        proto_send_ok(c->fd);

    } else if (u->console.type == CONSOLE_STDIN) {
        if (u->stdin_fd < 0) {
            proto_send_error(c->fd, "process not running");
            return;
        }
        ssize_t n = write(u->stdin_fd, data, strlen(data));
        if (n < 0) {
            proto_send_error(c->fd, "write to stdin failed");
            return;
        }
        proto_send_ok(c->fd);
    }
}

static void cmd_action(Client *c, const char *unit_name,
                       const char *action_name) {
    Unit *u = unit_find(unit_name);
    if (!u) { proto_send_error(c->fd, "unit not found"); return; }

    /* Find the action */
    const char *script = NULL;
    for (int i = 0; i < u->action_count; i++) {
        if (strcmp(u->actions[i].name, action_name) == 0) {
            script = u->actions[i].script;
            break;
        }
    }
    if (!script) { proto_send_error(c->fd, "action not found"); return; }

    log_info("Running action '%s' for unit %s: %s",
             action_name, unit_name, script);

    /* Fork a child to run the script */
    pid_t pid = fork();
    if (pid < 0) {
        proto_send_error(c->fd, "fork failed");
        return;
    }
    if (pid == 0) {
        /* Child: exec the script */
        execl("/bin/bash", "bash", "-c", script, NULL);
        _exit(127);
    }

    /* Parent: don't wait — SIGCHLD handler will reap it */
    proto_send_ok(c->fd);
}

/* ── Main dispatch ───────────────────────────────────────────────────────── */

int client_handle(int fd) {
    Client *c = client_find(fd);
    if (!c) return -1;

    char buf[PROTO_MAX_MSG];
    int n = proto_recv(fd, buf, sizeof(buf));

    if (n <= 0) return -1;  /* 0 = disconnect, -1 = error */

    /* Extract command and fields */
    char cmd[32] = {0};
    char unit[MAX_NAME] = {0};
    char action[MAX_NAME] = {0};
    char data[RING_BUF_LINE_MAX] = {0};

    json_get_str(buf, "cmd",     cmd,    sizeof(cmd));
    json_get_str(buf, "unit",    unit,   sizeof(unit));
    json_get_str(buf, "action",  action, sizeof(action));
    json_get_str(buf, "data",    data,   sizeof(data));

    /* For broadcast, "message" maps to the data field */
    if (strcmp(cmd, "broadcast") == 0)
        json_get_str(buf, "message", data, sizeof(data));

    log_debug("Client fd=%d cmd=%s unit=%s", fd, cmd, unit);

    if      (strcmp(cmd, "list")      == 0) cmd_list(c);
    else if (strcmp(cmd, "status")    == 0) cmd_status(c, unit);
    else if (strcmp(cmd, "tail")      == 0) cmd_tail(c, unit);
    else if (strcmp(cmd, "broadcast") == 0) cmd_broadcast(c, data);
    else if (strcmp(cmd, "attach")    == 0) cmd_attach(c, unit);
    else if (strcmp(cmd, "detach")    == 0) cmd_detach(c);
    else if (strcmp(cmd, "input")     == 0) cmd_input(c, unit, data);
    else if (strcmp(cmd, "action")    == 0) cmd_action(c, unit, action);
    else proto_send_error(fd, "unknown command");

    return 0;
}