/* 桥接客户端：与应用内的 ArkTS 服务通信（Unix socket 优先，TCP 兜底） */
#include "compat.h"

/* 复用一条连接：大块传输时每片都 connect/close 是白白多出来的开销 */
static int g_sock = -1;
static pthread_mutex_t g_lock = PTHREAD_MUTEX_INITIALIZER;

static int connect_tcp(void) {
    const char *host = getenv("USB_BRIDGE_HOST") ? getenv("USB_BRIDGE_HOST") : "127.0.0.1";
    int port = getenv("USB_BRIDGE_PORT") ? atoi(getenv("USB_BRIDGE_PORT")) : OHB_BRIDGE_PORT;
    int s = socket(AF_INET, SOCK_STREAM | SOCK_CLOEXEC, 0);
    if (s < 0) {
        return -1;
    }
    struct sockaddr_in sa;
    memset(&sa, 0, sizeof(sa));
    sa.sin_family = AF_INET;
    sa.sin_port = htons((unsigned short)port);
    if (inet_pton(AF_INET, host, &sa.sin_addr) != 1 ||
        connect(s, (struct sockaddr *)&sa, sizeof(sa)) != 0) {
        g_real_close(s);
        return -1;
    }
    return s;
}

static int connect_unix(void) {
    const char *path = getenv("USB_BRIDGE_SOCK") ? getenv("USB_BRIDGE_SOCK") : OHB_BRIDGE_SOCK;
    int s = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0);
    if (s < 0) {
        return -1;
    }
    struct sockaddr_un sa;
    memset(&sa, 0, sizeof(sa));
    sa.sun_family = AF_UNIX;
    snprintf(sa.sun_path, sizeof(sa.sun_path), "%s", path);
    if (connect(s, (struct sockaddr *)&sa, sizeof(sa)) != 0) {
        g_real_close(s);
        return -1;
    }
    return s;
}

int ohb_bridge_cmd(const char *cmd, char *reply, size_t replysz) {
    ohb_real_init();
    pthread_mutex_lock(&g_lock);
    int rc = -1;
    for (int attempt = 0; attempt < 2; attempt++) {
        if (g_sock < 0) {
            g_sock = connect_unix();
            if (g_sock < 0) {
                g_sock = connect_tcp();
            }
            if (g_sock < 0) {
                break;
            }
            struct timeval tv = { 5, 0 };
            setsockopt(g_sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
            setsockopt(g_sock, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv));
        }
        size_t len = strlen(cmd), off = 0;
        int ok = 1;
        while (off < len) {
            ssize_t n = send(g_sock, cmd + off, len - off, 0);
            if (n <= 0) {
                ok = 0;
                break;
            }
            off += (size_t)n;
        }
        if (ok && send(g_sock, "\n", 1, 0) != 1) {
            ok = 0;
        }
        size_t p = 0;
        if (ok) {
            while (p + 1 < replysz) {
                ssize_t n = recv(g_sock, reply + p, replysz - p - 1, 0);
                if (n <= 0) {
                    break;
                }
                p += (size_t)n;
                if (memchr(reply, '\n', p) != NULL) {
                    break;
                }
            }
        }
        reply[p] = '\0';
        char *nl = strchr(reply, '\n');
        if (nl) {
            *nl = '\0';
        }
        if (ok && p > 0) {
            rc = 0;
            break;
        }
        g_real_close(g_sock);
        g_sock = -1;
    }
    pthread_mutex_unlock(&g_lock);
    return rc;
}

/* ── 向 ArkTS 索取 usbfs fd ──
 *
 * 这是"数据面不过 ArkTS"的关键：拿到 fd 之后，claim/control/bulk 都是本进程
 * 对 fd 的普通 ioctl，桥接协议只在这一次握手时用到。
 * 失败返回 -1（老版本没有分发口，调用方按原逻辑退化为路径重写 + 桥接翻译）。
 */
/* 分发口里还没有该设备：先让 ArkTS 侧打开一次（会走 usbManager.requestRight
 * 并登记 fd），再重新索取。否则程序只能拿到自己 open 的节点，claim 必然失败。 */
static int ohb_acquire_fd_once(int bus, int addr);

static int kk_val(const char *reply, const char *key, int *out) {
    const char *p = strstr(reply, key);
    if (!p) {
        return 0;
    }
    return sscanf(p + strlen(key), "%d", out) == 1;
}

int ohb_bridge_open_ex(int bus, int addr, int *out_handle, int *out_bus, int *out_addr,
                       int *out_vid, int *out_pid) {
    char cmd[96];
    char reply[256];
    snprintf(cmd, sizeof(cmd), "open %d %d", bus, addr);
    int ok = ohb_bridge_cmd(cmd, reply, sizeof(reply)) == 0 && strncmp(reply, "OK", 2) == 0;
    if (!ok) {
        /* 旧地址已经不在了：用历史上同一地址对应的 VID/PID 反查新地址。
         * 设备重新枚举（拔插、系统重配）后这是唯一能自愈的途径。 */
        int vid = 0, pid = 0;
        if (!ohb_ident_lookup(bus, addr, &vid, &pid)) {
            return -1;
        }
        snprintf(cmd, sizeof(cmd), "openid %d %d", vid, pid);
        if (ohb_bridge_cmd(cmd, reply, sizeof(reply)) != 0 || strncmp(reply, "OK", 2) != 0) {
            return -1;
        }
        ohb_trace("open_ex %d-%d stale -> openid %d:%d", bus, addr, vid, pid);
    }
    int h = -1, nb = bus, na = addr, vid = 0, pid = 0;
    if (!kk_val(reply, "handle=", &h) || h < 0) {
        return -1;
    }
    kk_val(reply, "bus=", &nb);
    kk_val(reply, "addr=", &na);
    kk_val(reply, "vid=", &vid);
    kk_val(reply, "pid=", &pid);
    ohb_ident_remember(bus, addr, vid, pid);
    ohb_ident_remember(nb, na, vid, pid);
    if (out_handle) { *out_handle = h; }
    if (out_bus) { *out_bus = nb; }
    if (out_addr) { *out_addr = na; }
    if (out_vid) { *out_vid = vid; }
    if (out_pid) { *out_pid = pid; }
    return 0;
}

int ohb_acquire_fd(int bus, int addr) {
    int fd = ohb_acquire_fd_once(bus, addr);
    if (fd >= 0) {
        ohb_ident_learn_from_tree(bus, addr);
        return fd;
    }
    int h = -1, nb = bus, na = addr;
    if (ohb_bridge_open_ex(bus, addr, &h, &nb, &na, NULL, NULL) != 0) {
        return -1;
    }
    if (nb != bus || na != addr) {
        ohb_trace("acquire %d-%d relocated -> %d-%d", bus, addr, nb, na);
    }
    fd = ohb_acquire_fd_once(nb, na);
    if (fd >= 0) {
        ohb_ident_learn_from_tree(nb, na);
    }
    return fd;
}

static int ohb_acquire_fd_once(int bus, int addr) {
    const char *path = getenv("USB_FD_SOCK") ? getenv("USB_FD_SOCK") : OHB_FD_SOCK;
    int s = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0);
    if (s < 0) {
        return -1;
    }
    struct sockaddr_un sa;
    memset(&sa, 0, sizeof(sa));
    sa.sun_family = AF_UNIX;
    snprintf(sa.sun_path, sizeof(sa.sun_path), "%s", path);
    if (connect(s, (struct sockaddr *)&sa, sizeof(sa)) != 0) {
        g_real_close(s);
        return -1;
    }
    char req[64];
    int len = snprintf(req, sizeof(req), "get %d %d\n", bus, addr);
    if (send(s, req, (size_t)len, 0) != len) {
        g_real_close(s);
        return -1;
    }
    char status = 0;
    struct iovec iov;
    iov.iov_base = &status;
    iov.iov_len = 1;
    char ctrl[CMSG_SPACE(sizeof(int))];
    memset(ctrl, 0, sizeof(ctrl));
    struct msghdr msg;
    memset(&msg, 0, sizeof(msg));
    msg.msg_iov = &iov;
    msg.msg_iovlen = 1;
    msg.msg_control = ctrl;
    msg.msg_controllen = sizeof(ctrl);
    int fd = -1;
    if (recvmsg(s, &msg, 0) > 0 && status == 'F') {
        struct cmsghdr *cm = CMSG_FIRSTHDR(&msg);
        if (cm && cm->cmsg_level == SOL_SOCKET && cm->cmsg_type == SCM_RIGHTS) {
            memcpy(&fd, CMSG_DATA(cm), sizeof(int));
        }
    }
    g_real_close(s);
    if (fd >= 0) {
        /* 记住 bus/addr：read(fd) 时要按 usbfs 语义喂描述符链 */
        ohb_acq_register(fd, bus, addr);
    }
    return fd;
}

/* ── fork 安全：子进程不能继承一把已加锁的互斥量 ── */
void ohb_bridge_fork_prepare(void) {
    pthread_mutex_lock(&g_lock);
}

void ohb_bridge_fork_parent(void) {
    pthread_mutex_unlock(&g_lock);
}

void ohb_bridge_fork_child(void) {
    pthread_mutex_init(&g_lock, NULL);
    g_sock = -1;
}

/* ── base64（hex 在大块传输时两端都慢） ── */
static const char B64TAB[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

int ohb_b64_encode(const unsigned char *in, int n, char *out, int outsz) {
    int o = 0;
    int i = 0;
    for (; i + 2 < n && o + 5 < outsz; i += 3) {
        unsigned v = (unsigned)in[i] << 16 | (unsigned)in[i + 1] << 8 | in[i + 2];
        out[o++] = B64TAB[(v >> 18) & 63];
        out[o++] = B64TAB[(v >> 12) & 63];
        out[o++] = B64TAB[(v >> 6) & 63];
        out[o++] = B64TAB[v & 63];
    }
    int rem = n - i;
    if (rem > 0 && o + 5 < outsz) {
        unsigned v = (unsigned)in[i] << 16;
        if (rem == 2) {
            v |= (unsigned)in[i + 1] << 8;
        }
        out[o++] = B64TAB[(v >> 18) & 63];
        out[o++] = B64TAB[(v >> 12) & 63];
        out[o++] = rem == 2 ? B64TAB[(v >> 6) & 63] : '=';
        out[o++] = '=';
    }
    out[o] = '\0';
    return o;
}

static int b64_val(char c) {
    if (c >= 'A' && c <= 'Z') { return c - 'A'; }
    if (c >= 'a' && c <= 'z') { return c - 'a' + 26; }
    if (c >= '0' && c <= '9') { return c - '0' + 52; }
    if (c == '+') { return 62; }
    if (c == '/') { return 63; }
    return -1;
}

int ohb_b64_decode(const char *in, unsigned char *out, int max) {
    int o = 0;
    unsigned acc = 0;
    int bits = 0;
    for (const char *p = in; p && *p; p++) {
        int v = b64_val(*p);
        if (v < 0) {
            if (*p == '=') {
                break;
            }
            continue;
        }
        acc = (acc << 6) | (unsigned)v;
        bits += 6;
        if (bits >= 8) {
            bits -= 8;
            if (o < max) {
                out[o++] = (unsigned char)((acc >> bits) & 0xff);
            }
        }
    }
    return o;
}