#include "irc.h"

static void irc_parse_line(irc_client_t *c, const char *line);

/* ── base64 encode ─────────────────────────────────────────────────────── */
static const char B64[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
static void base64_encode(const unsigned char *in, size_t len, char *out, size_t out_sz) {
    size_t i = 0, j = 0;
    while (len > 0 && j + 5 < out_sz) {
        unsigned char a = in[i], b = len > 1 ? in[i+1] : 0, d = len > 2 ? in[i+2] : 0;
        out[j++] = B64[a >> 2];
        out[j++] = B64[((a & 3) << 4) | (b >> 4)];
        out[j++] = len > 1 ? B64[((b & 15) << 2) | (d >> 6)] : '=';
        out[j++] = len > 2 ? B64[d & 63] : '=';
        i += 3; len -= len >= 3 ? 3 : len;
    }
    out[j] = 0;
}

/* ── SASL helpers (called from receive thread) ─────────────────────────── */
static void sasl_send_registration(irc_client_t *c) {
    char line[256];
    snprintf(line, sizeof(line), "NICK %s", c->nick);
    irc_send_raw(c, line);
    snprintf(line, sizeof(line), "USER %s 0 * :A-Net IRC Client", c->nick);
    irc_send_raw(c, line);
}

static void sasl_send_credentials(irc_client_t *c) {
    /* PLAIN wire format: authzid\0authcid\0passwd  (authzid = empty) */
    unsigned char plain[384];
    int plen = 0;
    plain[plen++] = 0;
    int nl = (int)strlen(c->nick); memcpy(plain + plen, c->nick, nl); plen += nl;
    plain[plen++] = 0;
    int pl = (int)strlen(c->password); memcpy(plain + plen, c->password, pl); plen += pl;
    char b64[512]; base64_encode(plain, plen, b64, sizeof(b64));
    char out[560]; snprintf(out, sizeof(out), "AUTHENTICATE %s", b64);
    irc_send_raw(c, out);
}

static void sasl_finish(irc_client_t *c) {
    irc_send_raw(c, "CAP END");
    sasl_send_registration(c);
    c->sasl_state = SASL_DONE;
}

static void handle_cap(irc_client_t *c, const char *rest) {
    /* rest = "CAP <nick> <subcmd> :<caps>" */
    const char *s1 = strchr(rest, ' ');       /* skip "CAP" */
    if (!s1) return;
    const char *s2 = strchr(s1 + 1, ' ');    /* skip nick */
    if (!s2) return;
    const char *sub = s2 + 1;                 /* "LS :..." / "ACK :sasl" / "NAK :sasl" */

    if (starts_with(sub, "LS ")) {
        const char *caps = strchr(sub, ':');
        if (caps) caps++; else caps = sub + 3;
        if (c->sasl_state == SASL_WAIT_CAP) {
            if (strstr(caps, "sasl")) {
                irc_send_raw(c, "CAP REQ :sasl");
                c->sasl_state = SASL_WAIT_ACK;
            } else {
                irc_add_line(c, "*", "Server doesn't support SASL — connecting without auth", 1);
                sasl_finish(c);
            }
        } else {
            irc_send_raw(c, "CAP END");
        }
    } else if (starts_with(sub, "ACK")) {
        if (c->sasl_state == SASL_WAIT_ACK) {
            irc_send_raw(c, "AUTHENTICATE PLAIN");
            c->sasl_state = SASL_WAIT_PLUS;
        }
    } else if (starts_with(sub, "NAK")) {
        if (c->sasl_state != SASL_NONE && c->sasl_state != SASL_DONE) {
            irc_add_line(c, "*", "SASL rejected by server — connecting without auth", 1);
            sasl_finish(c);
        }
    }
}

void irc_init(irc_client_t *c) {
    memset(c, 0, sizeof(*c));
    c->sock = INVALID_SOCKET;
}

void irc_add_line(irc_client_t *c, const char *nick, const char *text, int status) {
    if (c->line_count >= UI_MAX_LINES) {
        memmove(&c->lines[0], &c->lines[1], sizeof(c->lines[0]) * (UI_MAX_LINES - 1));
        c->line_count = UI_MAX_LINES - 1;
    }
    chat_line_t *ln = &c->lines[c->line_count++];
    memset(ln, 0, sizeof(*ln));
    safe_strcpy(ln->nick, nick ? nick : "", sizeof(ln->nick));
    safe_strcpy(ln->text, text ? text : "", sizeof(ln->text));
    ln->ts = time(NULL);
    ln->status = status;
    c->dirty_chat = 1;
}

static int open_socket(const char *host, int port) {
    struct addrinfo hints, *res = NULL, *p;
    char portbuf[16];
    memset(&hints, 0, sizeof(hints));
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    snprintf(portbuf, sizeof(portbuf), "%d", port);
    if (getaddrinfo(host, portbuf, &hints, &res) != 0) return INVALID_SOCKET;
    socket_t s = INVALID_SOCKET;
    for (p = res; p; p = p->ai_next) {
        s = socket(p->ai_family, p->ai_socktype, p->ai_protocol);
        if (s == INVALID_SOCKET) continue;
        if (connect(s, p->ai_addr, (int)p->ai_addrlen) == 0) break;
        CLOSESOCK(s);
        s = INVALID_SOCKET;
    }
    freeaddrinfo(res);
    return s;
}

static int sock_read(socket_t s, char *buf, size_t len) {
#ifdef _WIN32
    return recv(s, buf, (int)len, 0);
#else
    return (int)recv(s, buf, len, 0);
#endif
}

int irc_send_raw(irc_client_t *c, const char *line) {
    if (!c->connected) return -1;
    char out[IRC_LINE_MAX + 4];
    size_t n = strlen(line);
    if (n > IRC_LINE_MAX - 2) n = IRC_LINE_MAX - 2;
    memcpy(out, line, n);
    out[n++] = '\r';
    out[n++] = '\n';
#ifdef _WIN32
    return send(c->sock, out, (int)n, 0) == (int)n ? 0 : -1;
#else
    return (int)send(c->sock, out, n, 0) == (int)n ? 0 : -1;
#endif
}

static void parse_names(irc_client_t *c, const char *names) {
    c->user_count = 0;
    char tmp[1024];
    safe_strcpy(tmp, names, sizeof(tmp));
    char *tok = strtok(tmp, " ");
    while (tok && c->user_count < UI_MAX_USERS) {
        while (*tok == '@' || *tok == '+' || *tok == '%') tok++;
        safe_strcpy(c->users[c->user_count++], tok, sizeof(c->users[0]));
        tok = strtok(NULL, " ");
    }
    c->dirty_users = 1;
}

static void handle_numeric(irc_client_t *c, int code, const char *line) {
    if (code == 1) {
        c->registered = 1;
        char msg[256];
        snprintf(msg, sizeof(msg), "Connected to %s:%d as %s", c->server, c->port, c->nick);
        irc_add_line(c, "*", msg, 1);
        if (c->channel[0]) {
            char joincmd[128];
            snprintf(joincmd, sizeof(joincmd), "JOIN %s", c->channel);
            irc_send_raw(c, joincmd);
        }
    } else if (code == 353) {
        const char *colon = strrchr(line, ':');
        if (colon && colon[1]) parse_names(c, colon + 1);
    } else if (code == 366) {
        c->dirty_users = 1;
    } else if (code == 372 || code == 375 || code == 376 || code == 477) {
        const char *colon = strstr(line, " :");
        irc_add_line(c, "*", colon ? colon + 2 : line, 1);
    }
}

static void irc_parse_line(irc_client_t *c, const char *line) {
    if (starts_with(line, "PING :")) {
        char pong[IRC_LINE_MAX];
        snprintf(pong, sizeof(pong), "PONG :%s", line + 6);
        irc_send_raw(c, pong);
        return;
    }

    /* SASL: server ready for credentials */
    if (starts_with(line, "AUTHENTICATE +")) {
        if (c->sasl_state == SASL_WAIT_PLUS) {
            sasl_send_credentials(c);
            c->sasl_state = SASL_WAIT_RESULT;
        }
        return;
    }

    if (line[0] == ':') {
        const char *sp = strchr(line, ' ');
        if (!sp) return;
        char prefix[128] = {0};
        snprintf(prefix, sizeof(prefix), "%.*s", (int)(sp - line - 1), line + 1);
        const char *rest = sp + 1;

        /* CAP negotiation */
        if (starts_with(rest, "CAP ")) {
            handle_cap(c, rest);
            return;
        }

        if (isdigit((unsigned char)rest[0]) && isdigit((unsigned char)rest[1]) && isdigit((unsigned char)rest[2])) {
            int code = atoi(rest);
            if (code == 903) {
                /* SASL authentication successful */
                irc_add_line(c, "*", "SASL authentication successful", 1);
                sasl_finish(c);
                return;
            }
            if (code == 904 || code == 905 || code == 906) {
                /* SASL failed / aborted */
                irc_add_line(c, "*", "SASL authentication failed — check password", 1);
                sasl_finish(c);
                return;
            }
            handle_numeric(c, code, line);
            return;
        }
        if (starts_with(rest, "PRIVMSG ")) {
            const char *msg = strstr(rest, " :");
            if (!msg) return;
            char nick[64] = {0};
            const char *bang = strchr(prefix, '!');
            snprintf(nick, sizeof(nick), "%.*s", bang ? (int)(bang - prefix) : (int)strlen(prefix), prefix);
            irc_add_line(c, nick, msg + 2, 0);
            return;
        }
        if (starts_with(rest, "JOIN ")) {
            char msg[128];
            snprintf(msg, sizeof(msg), "%s joined", prefix);
            irc_add_line(c, "*", msg, 1);
            irc_send_raw(c, "NAMES");
            return;
        }
        if (starts_with(rest, "PART ")) {
            char msg[128];
            snprintf(msg, sizeof(msg), "%s left", prefix);
            irc_add_line(c, "*", msg, 1);
            irc_send_raw(c, "NAMES");
            return;
        }
    }

    irc_add_line(c, "*", line, 1);
}

#ifdef _WIN32
static DWORD WINAPI irc_thread_main(LPVOID arg)
#else
static void *irc_thread_main(void *arg)
#endif
{
    irc_client_t *c = (irc_client_t *)arg;
    char buf[2048];
    while (c->thread_running && c->connected) {
        int r = sock_read(c->sock, buf, sizeof(buf) - 1);
        if (r <= 0) break;
        buf[r] = 0;
        for (int i = 0; i < r; i++) {
            char ch = buf[i];
            if (ch == '\n') {
                c->recvbuf[c->recv_used] = 0;
                trim(c->recvbuf);
                if (c->recvbuf[0]) irc_parse_line(c, c->recvbuf);
                c->recv_used = 0;
            } else if (ch != '\r' && c->recv_used < sizeof(c->recvbuf) - 1) {
                c->recvbuf[c->recv_used++] = ch;
            }
        }
    }
    c->connected = 0;
    c->thread_running = 0;
    c->dirty_status = 1;
    irc_add_line(c, "*", "Disconnected from server", 1);
#ifdef _WIN32
    return 0;
#else
    return NULL;
#endif
}

int irc_connect_server(irc_client_t *c, const char *server, int port,
                       const char *nick, const char *channel, int tls,
                       const char *password) {
    (void)tls;
    irc_disconnect(c, NULL);
    safe_strcpy(c->server, server, sizeof(c->server));
    safe_strcpy(c->nick, nick, sizeof(c->nick));
    safe_strcpy(c->channel, channel, sizeof(c->channel));
    safe_strcpy(c->password, password ? password : "", sizeof(c->password));
    c->port = port;
    c->tls = tls;
    c->sock = open_socket(server, port);
    if (c->sock == INVALID_SOCKET) {
        irc_add_line(c, "*", "Connection failed", 1);
        return -1;
    }
    c->connected = 1;
    c->thread_running = 1;
    c->recv_used = 0;
    irc_send_raw(c, "CAP LS 302");
    if (c->password[0]) {
        /* SASL path: defer NICK/USER until capability negotiation completes */
        c->sasl_state = SASL_WAIT_CAP;
    } else {
        /* No SASL: send NICK/USER immediately, end CAP negotiation via CAP END
           which the receive thread sends when it gets the LS response */
        c->sasl_state = SASL_NONE;
        sasl_send_registration(c);
    }
#ifdef _WIN32
    c->thread = CreateThread(NULL, 0, irc_thread_main, c, 0, NULL);
#else
    pthread_create(&c->thread, NULL, irc_thread_main, c);
#endif
    c->dirty_status = 1;
    return 0;
}

void irc_disconnect(irc_client_t *c, const char *reason) {
    if (!c->connected && c->sock == INVALID_SOCKET) return;
    if (c->connected) {
        char q[256];
        snprintf(q, sizeof(q), "QUIT :%s", reason ? reason : "Leaving");
        irc_send_raw(c, q);
    }
    c->connected = 0;
    if (c->sock != INVALID_SOCKET) {
        CLOSESOCK(c->sock);
        c->sock = INVALID_SOCKET;
    }
    c->thread_running = 0;
#ifndef _WIN32
    if (c->thread) pthread_join(c->thread, NULL);
#else
    if (c->thread) { WaitForSingleObject(c->thread, 500); CloseHandle(c->thread); c->thread = NULL; }
#endif
}

int irc_send_privmsg(irc_client_t *c, const char *target, const char *text) {
    char buf[IRC_LINE_MAX];
    snprintf(buf, sizeof(buf), "PRIVMSG %s :%s", target, text);
    if (irc_send_raw(c, buf) == 0) {
        irc_add_line(c, c->nick, text, 0);
        return 0;
    }
    return -1;
}

int irc_handle_command(irc_client_t *c, const char *input) {
    if (!input || !input[0]) return 0;
    if (input[0] != '/') {
        return irc_send_privmsg(c, c->channel[0] ? c->channel : c->nick, input);
    }
    char cmd[64];
    const char *args = strchr(input, ' ');
    int len = args ? (int)(args - input - 1) : (int)strlen(input + 1);
    snprintf(cmd, sizeof(cmd), "%.*s", len, input + 1);
    args = args ? args + 1 : "";
    if (strcasecmp(cmd, "join") == 0) {
        safe_strcpy(c->channel, args, sizeof(c->channel));
        char out[128]; snprintf(out, sizeof(out), "JOIN %s", args); return irc_send_raw(c, out);
    }
    if (strcasecmp(cmd, "part") == 0) {
        char out[128]; snprintf(out, sizeof(out), "PART %s", c->channel); return irc_send_raw(c, out);
    }
    if (strcasecmp(cmd, "nick") == 0) {
        safe_strcpy(c->nick, args, sizeof(c->nick));
        char out[128]; snprintf(out, sizeof(out), "NICK %s", args); return irc_send_raw(c, out);
    }
    if (strcasecmp(cmd, "msg") == 0) {
        char target[64] = {0};
        const char *sp = strchr(args, ' ');
        if (!sp) return -1;
        snprintf(target, sizeof(target), "%.*s", (int)(sp - args), args);
        return irc_send_privmsg(c, target, sp + 1);
    }
    if (strcasecmp(cmd, "raw") == 0 || strcasecmp(cmd, "quote") == 0) return irc_send_raw(c, args);
    if (strcasecmp(cmd, "who") == 0) return irc_send_raw(c, args[0] ? input + 1 : "WHO");
    if (strcasecmp(cmd, "names") == 0) return irc_send_raw(c, args[0] ? input + 1 : "NAMES");
    if (strcasecmp(cmd, "list") == 0) return irc_send_raw(c, "LIST");
    if (strcasecmp(cmd, "quit") == 0) { c->should_quit = 1; return 0; }
    if (strcasecmp(cmd, "help") == 0) {
        irc_add_line(c, "*", "Commands: /join /part /msg /nick /raw /who /names /list /quit", 1);
        return 0;
    }
    return irc_send_raw(c, input + 1);
}
