/*
 * umbrella-cli: command-line client for the umbrella daemon.
 *
 * Usage:
 *   umbrella-cli list
 *   umbrella-cli status <unit>
 *   umbrella-cli tail <unit>          # print recent output, then exit
 *   umbrella-cli attach <unit>        # interactive console session
 *   umbrella-cli input <unit> <cmd>   # send a single command
 *   umbrella-cli action <unit> <action>
 *   umbrella-cli broadcast <message>
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <fcntl.h>
#include <termios.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <sys/select.h>
#include <arpa/inet.h>

#define SOCK_PATH   "/run/umbrella/umbrella.sock"
#define PROTO_MAX   65536
#define HDR_SIZE    4

/* ── Connection ──────────────────────────────────────────────────────────── */

static int connect_daemon(void) {
    int fd = socket(AF_UNIX, SOCK_STREAM, 0);
    if (fd < 0) { perror("socket"); return -1; }

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

    if (connect(fd, (struct sockaddr *)&addr, sizeof(addr)) != 0) {
        if (errno == ENOENT)
            fprintf(stderr, "umbrella: daemon not running "
                    "(socket not found: %s)\n", SOCK_PATH);
        else
            perror("connect");
        close(fd);
        return -1;
    }
    return fd;
}

/* ── Protocol ────────────────────────────────────────────────────────────── */

static int send_msg(int fd, const char *json) {
    uint32_t len = (uint32_t)strlen(json);
    uint32_t net = htonl(len);
    if (write(fd, &net, 4) != 4) return -1;
    if (write(fd, json, len) != (ssize_t)len) return -1;
    return 0;
}

static int recv_msg(int fd, char *buf, int buf_size) {
    uint32_t net;
    int got = 0;
    while (got < 4) {
        int n = read(fd, (uint8_t*)&net + got, 4 - got);
        if (n <= 0) return -1;
        got += n;
    }
    uint32_t len = ntohl(net);
    if (!len || (int)len >= buf_size) return -1;
    got = 0;
    while (got < (int)len) {
        int n = read(fd, buf + got, len - got);
        if (n <= 0) return -1;
        got += n;
    }
    buf[len] = '\0';
    return (int)len;
}

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

static int jget(const char *json, const char *key, char *out, int size) {
    char needle[128];
    snprintf(needle, sizeof(needle), "\"%s\":", key);
    const char *p = strstr(json, needle);
    if (!p) return 0;
    p += strlen(needle);
    while (*p == ' ') p++;
    if (*p == '"') {
        p++;
        int i = 0;
        while (*p && *p != '"' && i < size - 1) {
		if (*p == '\\' && *(p+1)) {
			p++;
			switch (*p) {
				case 'n':	out[i++] = '\n'; break;
				case 'r':	out[i++] = '\r'; break;
				case 't':	out[i++] = '\t'; break;
				case '"':	out[i++] = '"'; break;
				default:	out[i++] = *p;	 break;
			}
			p++;
		} else {
			out[i++] = *p++;
		}
	}
        out[i] = '\0';
        return 1;
    } else {
        int i = 0;
        while (*p && *p != ',' && *p != '}' && i < size - 1)
            out[i++] = *p++;
        out[i] = '\0';
        return i > 0;
    }
}

/* ── Commands ────────────────────────────────────────────────────────────── */

static int cmd_list(int fd) {
    if (send_msg(fd, "{\"cmd\":\"list\"}") != 0) return 1;

    char buf[PROTO_MAX];
    if (recv_msg(fd, buf, sizeof(buf)) <= 0) return 1;

    printf("%-24s %-32s %-14s %-8s %s\n",
           "NAME", "DISPLAY", "STATE", "PLAYERS", "MAP");
    printf("%-24s %-32s %-14s %-8s %s\n",
           "────────────────────────",
           "────────────────────────────────",
           "──────────────",
           "────────",
           "───────────────────");

    const char *p = buf;
    while ((p = strstr(p, "\"name\":")) != NULL) {
        char name[64] = {0}, display[128] = {0}, state[32] = {0};
        char players_s[16] = {0}, max_players_s[16] = {0}, map[64] = {0};

        char snippet[512];
        strncpy(snippet, p, sizeof(snippet) - 1);

        jget(snippet, "name",        name,        sizeof(name));
        jget(snippet, "display",     display,     sizeof(display));
        jget(snippet, "state",       state,       sizeof(state));
        jget(snippet, "players",     players_s,   sizeof(players_s));
        jget(snippet, "max_players", max_players_s, sizeof(max_players_s));
        jget(snippet, "map",         map,         sizeof(map));

        if (!name[0]) { p++; continue; }

        /* Format player count as "N/M" or "-" */
        char players_fmt[16] = "-";
        int pl = players_s[0] ? atoi(players_s) : -1;
        int mx = max_players_s[0] ? atoi(max_players_s) : -1;
        if (pl >= 0 && mx >= 0)
            snprintf(players_fmt, sizeof(players_fmt), "%d/%d", pl, mx);

        printf("%-24s %-32s %-14s %-8s %s\n",
               name, display, state, players_fmt,
               map[0] ? map : "-");

        p++;
    }
    return 0;
}

static int cmd_status(int fd, const char *unit) {
    char msg[256];
    snprintf(msg, sizeof(msg),
             "{\"cmd\":\"status\",\"unit\":\"%s\"}", unit);
    if (send_msg(fd, msg) != 0) return 1;

    char buf[PROTO_MAX];
    if (recv_msg(fd, buf, sizeof(buf)) <= 0) return 1;

    char type[32] = {0};
    jget(buf, "type", type, sizeof(type));
    if (strcmp(type, "error") == 0) {
        char errmsg[256] = {0};
        jget(buf, "message", errmsg, sizeof(errmsg));
        fprintf(stderr, "Error: %s\n", errmsg);
        return 1;
    }

    char display[128] = {0}, state[32] = {0}, pid[16] = {0};
    char players_s[16] = {0}, max_players_s[16] = {0}, map[64] = {0};
    jget(buf, "display",     display,       sizeof(display));
    jget(buf, "state",       state,         sizeof(state));
    jget(buf, "pid",         pid,           sizeof(pid));
    jget(buf, "players",     players_s,     sizeof(players_s));
    jget(buf, "max_players", max_players_s, sizeof(max_players_s));
    jget(buf, "map",         map,           sizeof(map));

    printf("Unit    : %s\n", unit);
    printf("Display : %s\n", display);
    printf("State   : %s\n", state);
    if (pid[0] && strcmp(pid, "0") != 0)
        printf("PID     : %s\n", pid);

    int pl = players_s[0] ? atoi(players_s) : -1;
    int mx = max_players_s[0] ? atoi(max_players_s) : -1;
    if (pl >= 0 && mx >= 0)
        printf("Players : %d/%d\n", pl, mx);
    if (map[0])
        printf("Map     : %s\n", map);

    return 0;
}

static int cmd_tail(int fd, const char *unit) {
    char msg[256];
    snprintf(msg, sizeof(msg),
             "{\"cmd\":\"tail\",\"unit\":\"%s\"}", unit);
    if (send_msg(fd, msg) != 0) return 1;

    char buf[PROTO_MAX];
    if (recv_msg(fd, buf, sizeof(buf)) <= 0) return 1;

    char type[32] = {0};
    jget(buf, "type", type, sizeof(type));
    if (strcmp(type, "error") == 0) {
        char errmsg[256] = {0};
        jget(buf, "message", errmsg, sizeof(errmsg));
        fprintf(stderr, "Error: %s\n", errmsg);
        return 1;
    }

    char data[PROTO_MAX] = {0};
    jget(buf, "data", data, sizeof(data));
    if (data[0])
        printf("%s", data);
    return 0;
}

static int cmd_broadcast(int fd, const char *message) {
    char msg[1024];
    /* Simple JSON escaping for the message */
    char escaped[900] = {0};
    int j = 0;
    for (int i = 0; message[i] && j < (int)sizeof(escaped) - 2; i++) {
        if (message[i] == '"')       { escaped[j++] = '\\'; escaped[j++] = '"'; }
        else if (message[i] == '\\') { escaped[j++] = '\\'; escaped[j++] = '\\'; }
        else                          escaped[j++] = message[i];
    }
    escaped[j] = '\0';

    snprintf(msg, sizeof(msg), "{\"cmd\":\"broadcast\",\"message\":\"%s\"}", escaped);
    if (send_msg(fd, msg) != 0) return 1;

    char buf[PROTO_MAX];
    if (recv_msg(fd, buf, sizeof(buf)) <= 0) return 1;

    char type[32] = {0};
    jget(buf, "type", type, sizeof(type));
    if (strcmp(type, "error") == 0) {
        char errmsg[256] = {0};
        jget(buf, "message", errmsg, sizeof(errmsg));
        fprintf(stderr, "Error: %s\n", errmsg);
        return 1;
    }

    char sent_s[16] = {0}, failed_s[16] = {0};
    jget(buf, "sent",   sent_s,   sizeof(sent_s));
    jget(buf, "failed", failed_s, sizeof(failed_s));
    printf("Broadcast sent to %s unit(s), %s failed\n",
           sent_s[0] ? sent_s : "0",
           failed_s[0] ? failed_s : "0");
    return 0;
}

static int cmd_input(int fd, const char *unit, const char *input) {
    char msg[1024];
    snprintf(msg, sizeof(msg),
             "{\"cmd\":\"input\",\"unit\":\"%s\",\"data\":\"%s\\n\"}",
             unit, input);
    if (send_msg(fd, msg) != 0) return 1;

    char buf[PROTO_MAX];
    if (recv_msg(fd, buf, sizeof(buf)) <= 0) return 1;

    /* Print any output that came back */
    char type[32] = {0};
    jget(buf, "type", type, sizeof(type));
    if (strcmp(type, "output") == 0) {
        char data[4096] = {0};
        jget(buf, "data", data, sizeof(data));
        printf("%s\n", data);
    } else if (strcmp(type, "error") == 0) {
        char errmsg[256] = {0};
        jget(buf, "message", errmsg, sizeof(errmsg));
        fprintf(stderr, "Error: %s\n", errmsg);
        return 1;
    }
    return 0;
}

static int cmd_action(int fd, const char *unit, const char *action) {
    char msg[256];
    snprintf(msg, sizeof(msg),
             "{\"cmd\":\"action\",\"unit\":\"%s\",\"action\":\"%s\"}",
             unit, action);
    if (send_msg(fd, msg) != 0) return 1;

    char buf[PROTO_MAX];
    if (recv_msg(fd, buf, sizeof(buf)) <= 0) return 1;

    char type[32] = {0};
    jget(buf, "type", type, sizeof(type));
    if (strcmp(type, "error") == 0) {
        char errmsg[256] = {0};
        jget(buf, "message", errmsg, sizeof(errmsg));
        fprintf(stderr, "Error: %s\n", errmsg);
        return 1;
    }
    printf("Action '%s' dispatched for unit '%s'\n", action, unit);
    return 0;
}

/*
 * cmd_attach: interactive console session.
 * Attaches to the unit, receives live output, sends typed input.
 * Ctrl+D or Ctrl+C to detach.
 */
static int cmd_attach(int fd, const char *unit) {
    char msg[256];
    snprintf(msg, sizeof(msg),
             "{\"cmd\":\"attach\",\"unit\":\"%s\"}", unit);
    if (send_msg(fd, msg) != 0) return 1;

    /* Set terminal to raw mode so we get input character by character */
    struct termios orig, raw;
    int is_tty = isatty(STDIN_FILENO);
    if (is_tty) {
        tcgetattr(STDIN_FILENO, &orig);
        raw = orig;
        raw.c_lflag &= ~(ICANON | ECHO);
        raw.c_cc[VMIN]  = 1;
        raw.c_cc[VTIME] = 0;
        tcsetattr(STDIN_FILENO, TCSANOW, &raw);
    }

    printf("── Attached to %s (Ctrl+D to detach) ──\n", unit);
    char drain[PROTO_MAX];
    recv_msg(fd, drain, sizeof(drain));
    fflush(stdout);

    char input_buf[1024];
    int input_pos = 0;
    int ret = 0;

    while (1) {
        fd_set rfds;
        FD_ZERO(&rfds);
        FD_SET(fd, &rfds);          /* daemon messages */
        FD_SET(STDIN_FILENO, &rfds); /* keyboard input */
        int maxfd = fd > STDIN_FILENO ? fd : STDIN_FILENO;

        int n = select(maxfd + 1, &rfds, NULL, NULL, NULL);
        if (n < 0) {
            if (errno == EINTR) continue;
            break;
        }

        /* Daemon sent us something */
        if (FD_ISSET(fd, &rfds)) {
            char buf[PROTO_MAX];
            int r = recv_msg(fd, buf, sizeof(buf));
            if (r <= 0) break;

            char type[32] = {0};
            jget(buf, "type", type, sizeof(type));

            if (strcmp(type, "output") == 0) {
                char data[PROTO_MAX] = {0};
                jget(buf, "data", data, sizeof(data));
                /* Print output, restore cursor line for input */
                printf("%s", data);
                fflush(stdout);
            } else if (strcmp(type, "error") == 0) {
                char errmsg[256] = {0};
                jget(buf, "message", errmsg, sizeof(errmsg));
                fprintf(stderr, "\nError: %s\n", errmsg);
                ret = 1;
                break;
            }
        }

        /* User typed something */
        if (FD_ISSET(STDIN_FILENO, &rfds)) {
            char ch;
            ssize_t r = read(STDIN_FILENO, &ch, 1);
            if (r <= 0) break;

            if (ch == 4) break; /* Ctrl+D — detach */

            if (ch == '\n' || ch == '\r') {
                /* Send the buffered line to the unit */
                input_buf[input_pos] = '\0';
                printf("\n");

                char send_buf[1200];
                /* JSON-escape the input */
                char escaped[1100];
                int j = 0;
                for (int i = 0; input_buf[i] && j < 1090; i++) {
                    if (input_buf[i] == '"')  { escaped[j++] = '\\'; escaped[j++] = '"'; }
                    else if (input_buf[i] == '\\') { escaped[j++] = '\\'; escaped[j++] = '\\'; }
                    else escaped[j++] = input_buf[i];
                }
                escaped[j] = '\0';

                snprintf(send_buf, sizeof(send_buf),
                         "{\"cmd\":\"input\",\"unit\":\"%s\","
                         "\"data\":\"%s\\n\"}", unit, escaped);
                send_msg(fd, send_buf);
                input_pos = 0;
            } else if (ch == 127 || ch == '\b') {
                /* Backspace */
                if (input_pos > 0) {
                    input_pos--;
                    printf("\b \b");
                    fflush(stdout);
                }
            } else {
                if (input_pos < (int)sizeof(input_buf) - 1) {
                    input_buf[input_pos++] = ch;
                    /* Echo the character */
                    write(STDOUT_FILENO, &ch, 1);
                }
            }
        }
    }

    /* Restore terminal */
    if (is_tty)
        tcsetattr(STDIN_FILENO, TCSANOW, &orig);

    /* Send detach */
    send_msg(fd, "{\"cmd\":\"detach\"}");
    printf("\n── Detached ──\n");
    return ret;
}

/* ── Entry point ─────────────────────────────────────────────────────────── */

static void usage(void) {
    fprintf(stderr,
            "Usage: umbrella-cli <command> [args]\n"
            "\n"
            "Commands:\n"
            "  list                          List all units\n"
            "  status <unit>                 Show unit status\n"
            "  tail <unit>                   Print recent output and exit\n"
            "  attach <unit>                 Interactive console session\n"
            "  input  <unit> <cmd>           Send a single command\n"
            "  action <unit> <action>        Run a named action\n"
            "  broadcast <message>           Send message to all running units\n"
            "\n"
            "Examples:\n"
            "  umbrella-cli list\n"
            "  umbrella-cli tail tf2-novemen\n"
            "  umbrella-cli attach tf2-novemen\n"
            "  umbrella-cli input tf2-novemen \"say Server restarting soon\"\n"
            "  umbrella-cli action tf2-novemen update\n"
            "  umbrella-cli broadcast \"Server restart in 5 minutes\"\n");
}

int main(int argc, char *argv[]) {
    if (argc < 2) { usage(); return 1; }

    int fd = connect_daemon();
    if (fd < 0) return 1;

    int ret = 0;
    const char *cmd = argv[1];

    if (strcmp(cmd, "list") == 0) {
        ret = cmd_list(fd);
    } else if (strcmp(cmd, "status") == 0) {
        if (argc < 3) { fprintf(stderr, "status: need unit name\n"); ret = 1; }
        else ret = cmd_status(fd, argv[2]);
    } else if (strcmp(cmd, "tail") == 0) {
        if (argc < 3) { fprintf(stderr, "tail: need unit name\n"); ret = 1; }
        else ret = cmd_tail(fd, argv[2]);
    } else if (strcmp(cmd, "attach") == 0) {
        if (argc < 3) { fprintf(stderr, "attach: need unit name\n"); ret = 1; }
        else ret = cmd_attach(fd, argv[2]);
    } else if (strcmp(cmd, "input") == 0) {
        if (argc < 4) { fprintf(stderr, "input: need unit and command\n"); ret = 1; }
        else ret = cmd_input(fd, argv[2], argv[3]);
    } else if (strcmp(cmd, "action") == 0) {
        if (argc < 4) { fprintf(stderr, "action: need unit and action name\n"); ret = 1; }
        else ret = cmd_action(fd, argv[2], argv[3]);
    } else if (strcmp(cmd, "broadcast") == 0) {
        if (argc < 3) { fprintf(stderr, "broadcast: need message\n"); ret = 1; }
        else ret = cmd_broadcast(fd, argv[2]);
    } else {
        fprintf(stderr, "Unknown command: %s\n\n", cmd);
        usage();
        ret = 1;
    }

    close(fd);
    return ret;
}