#include "plc.h"
#include "flog.h"

#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <poll.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <stdarg.h>
#include <unistd.h>

#define S7_PDU_REQ    0x01
#define S7_PDU_ACKD   0x03

#define COTP_CR       0xE0
#define COTP_CC       0xD0
#define COTP_DT       0xF0

#define TS_BYTE       0x02

#define SOCK_TIMEOUT_MS 3000

typedef struct {
    uint8_t buf[1024];
    int len;
} Frame;

typedef struct {
    uint8_t data[PLC_MAX_BLOCK];
    int     len;
} PlcBlock;

static PlcConfig g_cfg;
static int g_inited = 0;
static pthread_t g_thread;
static int g_running = 0;
static pthread_mutex_t g_mutex = PTHREAD_MUTEX_INITIALIZER;

static int g_sock = -1;
static int g_connected = 0;
static PlcBlock g_block;
static uint16_t g_seq = 0;
static char g_err[256];

#define DBG(...) do { if (getenv("PLC_DEBUG")) fprintf(stderr, __VA_ARGS__); } while (0)

static void set_err(const char *fmt, ...) {
    va_list ap;
    va_start(ap, fmt);
    vsnprintf(g_err, sizeof(g_err), fmt, ap);
    va_end(ap);
}

static int recv_exact(int fd, void *buf, size_t len) {
    size_t got = 0;
    while (got < len) {
        ssize_t r = recv(fd, (uint8_t *)buf + got, len - got, 0);
        if (r <= 0) {
            if (r < 0 && (errno == EAGAIN || errno == EWOULDBLOCK))
                set_err("recv timeout");
            else
                set_err("recv failed: %s", strerror(errno));
            return -1;
        }
        got += (size_t)r;
    }
    return 0;
}

static int send_all(int fd, const void *buf, size_t len) {
    size_t sent = 0;
    while (sent < len) {
        ssize_t r = send(fd, (const uint8_t *)buf + sent, len - sent, MSG_NOSIGNAL);
        if (r <= 0) {
            if (r < 0 && (errno == EAGAIN || errno == EWOULDBLOCK))
                set_err("send timeout");
            else
                set_err("send failed: %s", strerror(errno));
            return -1;
        }
        sent += (size_t)r;
    }
    return 0;
}

/* Wrap payload in TPKT + COTP-DT and send. */
static int send_s7(int fd, const uint8_t *s7, int s7len) {
    uint8_t frame[4 + 3 + 1024];
    if (3 + s7len > (int)sizeof(frame) - 4) {
        set_err("s7 frame too large");
        return -1;
    }
    int total = 4 + 3 + s7len;
    frame[0] = 3;
    frame[1] = 0;
    frame[2] = (uint8_t)(total >> 8);
    frame[3] = (uint8_t)(total & 0xFF);
    frame[4] = 2;          /* COTP DT header length */
    frame[5] = COTP_DT;
    frame[6] = 0x80;       /* EOT + sequence 0 */
    memcpy(frame + 7, s7, s7len);
    return send_all(fd, frame, total);
}

/* Receive one TPKT frame, return COTP payload (sans 3-byte DT header). */
static int recv_s7(int fd, uint8_t *out, int maxout, int *out_len) {
    uint8_t hdr[4];
    if (recv_exact(fd, hdr, 4) < 0) return -1;
    if (hdr[0] != 3) {
        set_err("bad TPKT version %d", hdr[0]);
        return -1;
    }
    int total = ((int)hdr[2] << 8) | hdr[3];
    if (total < 7 || total > 4 + 1024) {
        set_err("bad TPKT length %d", total);
        return -1;
    }
    uint8_t payload[1024];
    if (recv_exact(fd, payload, total - 4) < 0) return -1;
    if (payload[0] != 2 || payload[1] != COTP_DT) {
        set_err("expected COTP DT, got %02x %02x", payload[0], payload[1]);
        return -1;
    }
    int s7len = total - 4 - 3;
    if (s7len < 0 || s7len > maxout) {
        set_err("bad s7 length %d", s7len);
        return -1;
    }
    memcpy(out, payload + 3, s7len);
    *out_len = s7len;
    return 0;
}

static void put_u16(uint8_t *p, uint16_t v) {
    p[0] = (uint8_t)(v >> 8);
    p[1] = (uint8_t)(v & 0xFF);
}

static uint16_t get_u16(const uint8_t *p) {
    return (uint16_t)(((uint16_t)p[0] << 8) | p[1]);
}

/* ---- COTP Connection Request ---- */
static int cotp_connect(int fd) {
    uint8_t pdu[22];
    int n = 0;
    pdu[n++] = 0x11;          /* length (excl. this byte) */
    pdu[n++] = COTP_CR;
    pdu[n++] = 0x00; pdu[n++] = 0x00;   /* dst ref */
    pdu[n++] = 0x00; pdu[n++] = 0x01;   /* src ref */
    pdu[n++] = 0x00;                    /* class */
    pdu[n++] = 0xC1; pdu[n++] = 0x02;   /* calling TSAP (local) */
    pdu[n++] = (uint8_t)((g_cfg.tsap_local >> 8) & 0xFF);
    pdu[n++] = (uint8_t)(g_cfg.tsap_local & 0xFF);
    pdu[n++] = 0xC2; pdu[n++] = 0x02;   /* called TSAP (remote) */
    pdu[n++] = (uint8_t)((g_cfg.tsap_remote >> 8) & 0xFF);
    pdu[n++] = (uint8_t)(g_cfg.tsap_remote & 0xFF);
    pdu[n++] = 0xC0; pdu[n++] = 0x01; pdu[n++] = 0x0A; /* PDU size 1024 */

    uint8_t frame[4 + n];
    int total = 4 + n;
    frame[0] = 3; frame[1] = 0;
    frame[2] = (uint8_t)(total >> 8);
    frame[3] = (uint8_t)(total & 0xFF);
    memcpy(frame + 4, pdu, n);
    if (send_all(fd, frame, total) < 0) return -1;

    fprintf(stderr, "[plc] COTP CR (%d bytes): ", total);
    for (int i = 0; i < total; i++) fprintf(stderr, "%02x ", frame[i]);
    fprintf(stderr, "\n");

    uint8_t resp[64];
    int rlen;
    if (recv_exact(fd, resp, 4) < 0) {
        fprintf(stderr, "[plc] no CC response (reset after CR)\n");
        return -1;
    }
    if (resp[0] != 3) { set_err("CC bad TPKT version"); return -1; }
    int rtot = ((int)resp[2] << 8) | resp[3];
    if (rtot < 7) { set_err("CC too short"); return -1; }
    rlen = rtot - 4;
    if (recv_exact(fd, resp + 4, rlen) < 0) return -1;
    if (resp[4] != 0x0B /* LI=11 */ || resp[5] != COTP_CC) {
        /* Some stacks vary LI; only type must be CC. */
        if (resp[5] != COTP_CC) { set_err("expected COTP CC"); return -1; }
    }
    DBG("[plc] COTP connected (CC)\n");
    return 0;
}

/* ---- S7 Setup Communication (negotiate PDU length) ---- */
static int s7_setup(int fd) {
    uint8_t req[18];
    req[0] = 0x32; req[1] = S7_PDU_REQ;
    req[2] = 0x00; req[3] = 0x00;
    put_u16(req + 4, ++g_seq);
    put_u16(req + 6, 0x0008);   /* param len */
    put_u16(req + 8, 0x0000);   /* data len */
    req[10] = 0xF0; req[11] = 0x00;   /* setup communication */
    put_u16(req + 12, 0x0001);  /* max AMQ caller */
    put_u16(req + 14, 0x0001);  /* max AMQ callee */
    put_u16(req + 16, 0x01E0);  /* pdu length 480 */

    if (send_s7(fd, req, sizeof(req)) < 0) return -1;

    uint8_t resp[1024];
    int rlen;
    if (recv_s7(fd, resp, sizeof(resp), &rlen) < 0) return -1;
    if (rlen < 12 || resp[0] != 0x32) {
        set_err("bad setup response");
        return -1;
    }
    /* Tolerate an error response; negotiation is non-critical. */
    DBG("[plc] setup done (%d bytes)\n", rlen);
    return 0;
}

/* ---- S7 Read Request ---- */
static int s7_read(int fd, uint8_t *block, int len) {
    uint8_t req[10 + 14];
    req[0] = 0x32; req[1] = S7_PDU_REQ;
    req[2] = 0x00; req[3] = 0x00;
    put_u16(req + 4, ++g_seq);
    put_u16(req + 6, 0x000E);   /* param len */
    put_u16(req + 8, 0x0000);   /* data len */
    req[10] = 0x04; req[11] = 0x01; /* READ VAR, 1 item */
    req[12] = 0x12;             /* spec */
    req[13] = 0x0A;             /* addr spec length */
    req[14] = 0x10;             /* syntax S7-ANY */
    req[15] = TS_BYTE;          /* transport size: BYTE */
    put_u16(req + 16, (uint16_t)len);
    put_u16(req + 18, g_cfg.area == PLC_AREA_DB ? (uint16_t)g_cfg.db : 0);
    req[20] = (uint8_t)g_cfg.area;
    uint32_t addr = (uint32_t)g_cfg.read_start * 8u;
    req[21] = (uint8_t)(addr >> 16);
    req[22] = (uint8_t)(addr >> 8);
    req[23] = (uint8_t)(addr & 0xFF);

    if (send_s7(fd, req, sizeof(req)) < 0) return -1;

    uint8_t resp[1024];
    int rlen;
    if (recv_s7(fd, resp, sizeof(resp), &rlen) < 0) return -1;
    if (rlen < 14 || resp[0] != 0x32) {
        set_err("bad read response header");
        return -1;
    }
    int pdu_type = resp[1];
    if (pdu_type != S7_PDU_ACKD) {
        set_err("unexpected read pdu type %02x", pdu_type);
        return -1;
    }
    uint16_t seq = get_u16(resp + 4);
    uint16_t param_len = get_u16(resp + 6);
    uint16_t data_len = get_u16(resp + 8);
    uint8_t err_cls = resp[10], err_code = resp[11];
    if (err_cls != 0) {
        set_err("s7 error class %02x code %02x", err_cls, err_code);
        return -1;
    }
    if (seq != g_seq) {
        /* Real S7 PLCs echo the request sequence; some simulators reply
         * with their own counter. Log but do not fail. */
        DBG("[plc] seq warn: resp %u != sent %u\n", seq, g_seq);
    }
    int off = 12 + param_len;
    int avail = rlen - off;
    if (avail < 4) {
        set_err("read response truncated");
        return -1;
    }
    int dlen = data_len;
    if (dlen > avail) {
        DBG("[plc] data_len warn: field %u > avail %d, clamping\n",
            data_len, avail);
        dlen = avail;
    }
    const uint8_t *d = resp + off;
    /* item: rc(1) ts(1) bitlen(2) data */
    if (dlen < 4) { set_err("read data empty"); return -1; }
    uint8_t rc = d[0];
    uint8_t ts = d[1];
    uint16_t blen = get_u16(d + 2);
    if (rc != 0xFF) {
        set_err("read item rc=0x%02x", rc);
        return -1;
    }
    int nbytes = (ts == TS_BYTE) ? blen : blen / 8;
    if (nbytes > (int)data_len - 4) {
        set_err("read data length mismatch (%d > %d)", nbytes, data_len - 4);
        return -1;
    }
    if (nbytes > PLC_MAX_BLOCK) nbytes = PLC_MAX_BLOCK;
    memcpy(block, d + 4, nbytes);
    return nbytes;
}

static int tcp_connect_handshake(void) {
    int fd = socket(AF_INET, SOCK_STREAM, 0);
    if (fd < 0) { set_err("socket: %s", strerror(errno)); return -1; }

    struct sockaddr_in sa;
    memset(&sa, 0, sizeof(sa));
    sa.sin_family = AF_INET;
    sa.sin_port = htons((uint16_t)g_cfg.port);
    if (inet_pton(AF_INET, g_cfg.ip, &sa.sin_addr) != 1) {
        set_err("bad ip %s", g_cfg.ip);
        close(fd);
        return -1;
    }

    /* Non-blocking connect + poll, so a slow/unreachable PLC cannot stall
     * the reconnect loop. */
    int flags = fcntl(fd, F_GETFL, 0);
    fcntl(fd, F_SETFL, flags | O_NONBLOCK);
    int rc = connect(fd, (struct sockaddr *)&sa, sizeof(sa));
    if (rc < 0 && errno != EINPROGRESS) {
        set_err("tcp connect: %s", strerror(errno));
        close(fd);
        return -1;
    }
    if (rc < 0) {
        struct pollfd pfd = { .fd = fd, .events = POLLOUT };
        int pr = poll(&pfd, 1, SOCK_TIMEOUT_MS);
        if (pr <= 0) {
            set_err(pr == 0 ? "tcp connect timeout" : "tcp connect: %s",
                    strerror(errno));
            close(fd);
            return -1;
        }
        int soerr = 0;
        socklen_t sl = sizeof(soerr);
        if (getsockopt(fd, SOL_SOCKET, SO_ERROR, &soerr, &sl) < 0 || soerr != 0) {
            set_err("tcp connect refused: %s", strerror(soerr));
            close(fd);
            return -1;
        }
    }
    fcntl(fd, F_SETFL, flags);

    struct timeval tv;
    tv.tv_sec = SOCK_TIMEOUT_MS / 1000;
    tv.tv_usec = (SOCK_TIMEOUT_MS % 1000) * 1000;
    setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
    setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv));
    int one = 1;
    setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one));

    if (cotp_connect(fd) < 0) { close(fd); return -1; }
    if (s7_setup(fd) < 0) {
        /* Setup is not fatal; keep the transport connection. */
        DBG("[plc] setup failed (%s), continuing\n", plc_last_error());
    }
    g_sock = fd;
    return 0;
}

static void *plc_thread(void *arg) {
    (void)arg;
    int was_err_logged = 0;
    while (g_running) {
        if (g_sock < 0) {
            if (tcp_connect_handshake() < 0) {
                if (!was_err_logged) {
                    fprintf(stderr, "[plc] connect %s:%d failed: %s\n",
                            g_cfg.ip, g_cfg.port, plc_last_error());
                    flog("PLC 连接 %s:%d 失败: %s", g_cfg.ip, g_cfg.port, plc_last_error());
                    was_err_logged = 1;
                }
                usleep((useconds_t)g_cfg.reconnect_ms * 1000);
                continue;
            }
            was_err_logged = 0;
            DBG("[plc] connected to %s:%d\n", g_cfg.ip, g_cfg.port);
        }

        uint8_t block[PLC_MAX_BLOCK];
        int n = s7_read(g_sock, block, g_cfg.read_len);
        if (n < 0) {
            fprintf(stderr, "[plc] read failed: %s\n", plc_last_error());
            flog_throttled("plc_read", 10000, "PLC 读取失败: %s", plc_last_error());
            close(g_sock);
            g_sock = -1;
            pthread_mutex_lock(&g_mutex);
            g_connected = 0;
            pthread_mutex_unlock(&g_mutex);
            usleep((useconds_t)g_cfg.reconnect_ms * 1000);
            continue;
        }

        pthread_mutex_lock(&g_mutex);
        memcpy(g_block.data, block, n);
        g_block.len = n;
        g_connected = 1;
        pthread_mutex_unlock(&g_mutex);

        DBG("[plc] read %d bytes ok\n", n);
        usleep((useconds_t)g_cfg.poll_ms * 1000);
    }
    return NULL;
}

int plc_init(const PlcConfig *cfg) {
    if (g_inited) return 0;
    if (!cfg || cfg->port <= 0 || cfg->read_len <= 0 ||
        cfg->read_len > PLC_MAX_BLOCK) {
        set_err("invalid plc config");
        return -1;
    }
    g_cfg = *cfg;
    g_cfg.ip[sizeof(g_cfg.ip) - 1] = '\0';
    if (g_cfg.poll_ms <= 0) g_cfg.poll_ms = 100;
    if (g_cfg.reconnect_ms <= 0) g_cfg.reconnect_ms = 3000;

    g_running = 1;
    g_sock = -1;
    g_connected = 0;
    if (pthread_create(&g_thread, NULL, plc_thread, NULL) != 0) {
        g_running = 0;
        set_err("pthread_create failed");
        return -1;
    }
    g_inited = 1;
    return 0;
}

void plc_shutdown(void) {
    if (!g_inited) return;
    g_running = 0;
    if (g_sock >= 0) {
        shutdown(g_sock, SHUT_RDWR);
        close(g_sock);
        g_sock = -1;
    }
    pthread_join(g_thread, NULL);
    g_inited = 0;
}

int plc_is_connected(void) {
    pthread_mutex_lock(&g_mutex);
    int c = g_connected;
    pthread_mutex_unlock(&g_mutex);
    return c;
}

int plc_get_bit(int byte_idx, int bit_idx) {
    int r = -1;
    pthread_mutex_lock(&g_mutex);
    if (g_connected && byte_idx >= 0 && byte_idx < g_block.len &&
        bit_idx >= 0 && bit_idx < 8) {
        r = (g_block.data[byte_idx] >> bit_idx) & 1;
    }
    pthread_mutex_unlock(&g_mutex);
    return r;
}

int plc_get_byte(int byte_idx, uint8_t *out) {
    int r = -1;
    pthread_mutex_lock(&g_mutex);
    if (g_connected && out && byte_idx >= 0 && byte_idx < g_block.len) {
        *out = g_block.data[byte_idx];
        r = 0;
    }
    pthread_mutex_unlock(&g_mutex);
    return r;
}

const char *plc_last_error(void) {
    pthread_mutex_lock(&g_mutex);
    if (g_err[0] == '\0') {
        pthread_mutex_unlock(&g_mutex);
        return "no error";
    }
    static char copy[256];
    strncpy(copy, g_err, sizeof(copy) - 1);
    copy[sizeof(copy) - 1] = '\0';
    pthread_mutex_unlock(&g_mutex);
    return copy;
}