#include "dm_identity.h"
#include "dm_crypto.h"
#include "../utun_instance.h"
#include "../routing_layer/devgroup.h"
#include "../transport_layer/secure_channel.h"
#include "../../lib/debug_config.h"
#include <sqlite3.h>
#include <openssl/crypto.h>
#include <openssl/sha.h>
#include <string.h>
#include <stdlib.h>

static int identity_error(struct UTUN_INSTANCE* inst, const char* stage, uint64_t node) {
    DEBUG_ERROR(DEBUG_CATEGORY_DM, "DM identity: %s node=%llu sql=%s", stage, (unsigned long long)node,
        inst && inst->topo_sqlite_db ? sqlite3_errmsg(inst->topo_sqlite_db) : "unavailable");
    return -1;
}

/* Проверка формы не зависит от доступности устройства и локального каталога. */
static int party_valid(const struct dm_party* party) {
    static const uint8_t zero[32] = {0};
    if (!party || !memcmp(party->ed, zero, 32) || !memcmp(party->x, zero, 32)) return 0;
    if (party->kind == DM_PARTY_GROUP) return party->node == 0;
    return party->kind == DM_PARTY_DEVICE && party->node && party->node == sc_derive_node_id_from_pubkey(party->x);
}

int dm_party_encode(const struct dm_party* party, uint8_t out[DM_PARTY_WIRE_SIZE]) {
    if (!out || !party_valid(party)) return identity_error(NULL, "invalid party encode", party ? party->node : 0);
    out[0] = party->kind;
    for (unsigned i = 0; i < 8; i++) out[1 + i] = (uint8_t)(party->node >> (8 * i));
    memcpy(out + 9, party->ed, 32); memcpy(out + 41, party->x, 32);
    return 0;
}

int dm_party_decode(const uint8_t data[DM_PARTY_WIRE_SIZE], struct dm_party* out) {
    if (!data || !out) return identity_error(NULL, "invalid party decode arguments", 0);
    struct dm_party party = {.kind = data[0]};
    for (unsigned i = 0; i < 8; i++) party.node |= (uint64_t)data[1 + i] << (8 * i);
    memcpy(party.ed, data + 9, 32); memcpy(party.x, data + 41, 32);
    if (!party_valid(&party)) return identity_error(NULL, "invalid party encoding", party.node);
    *out = party;
    return 0;
}

int dm_party_equal(const struct dm_party* a, const struct dm_party* b) {
    return a && b && a->kind == b->kind && a->node == b->node && !memcmp(a->ed, b->ed, 32) && !memcmp(a->x, b->x, 32);
}

/* Собственное устройство берём из ядра, удалённое — из проверенной проекции nodes. */
int dm_party_resolve(struct UTUN_INSTANCE* inst, uint64_t node, enum dm_mode mode, struct dm_party* out) {
    if (!inst || !inst->topo_sqlite_db || !node || !out || (mode != DM_MODE_DEVICE && mode != DM_MODE_GROUPS))
        return identity_error(inst, "invalid resolve arguments", node);
    struct dm_party party = {.kind = DM_PARTY_DEVICE, .node = node};
    if (node == inst->node_id) {
        memcpy(party.ed, inst->my_ed25519_pubkey, 32); memcpy(party.x, inst->my_keys.public_key, 32);
        if (mode == DM_MODE_GROUPS) {
            struct devgroup_record record; uint8_t ed[32] = {0}, x[32] = {0};
            int found = devgroup_load_local(inst, &record, ed, x);
            OPENSSL_cleanse(ed, sizeof(ed)); OPENSSL_cleanse(x, sizeof(x));
            if (found < 0) return identity_error(inst, "local subgroup verification failed", node);
            if (found) {
                party.kind = DM_PARTY_GROUP; party.node = 0;
                memcpy(party.ed, record.ed, 32); memcpy(party.x, record.x, 32);
            }
        }
    } else {
        sqlite3_stmt* st = NULL;
        int rc = sqlite3_prepare_v2(inst->topo_sqlite_db,
            "SELECT x25519_pubkey,ed25519_pubkey FROM nodes WHERE node_id=?", -1, &st, NULL);
        if (rc == SQLITE_OK) rc = sqlite3_bind_int64(st, 1, (sqlite3_int64)node);
        if (rc == SQLITE_OK) rc = sqlite3_step(st);
        if (rc == SQLITE_ROW && sqlite3_column_bytes(st, 0) == 32 && sqlite3_column_bytes(st, 1) == 32) {
            memcpy(party.x, sqlite3_column_blob(st, 0), 32); memcpy(party.ed, sqlite3_column_blob(st, 1), 32);
        } else {
            sqlite3_finalize(st);
            return identity_error(inst, "missing or damaged device keys", node);
        }
        if (sqlite3_finalize(st) != SQLITE_OK) return identity_error(inst, "device query finalization failed", node);
        if (!party_valid(&party)) return identity_error(inst, "device ID/key mismatch", node);
        if (mode == DM_MODE_GROUPS) {
            uint8_t id[32]; char name[64];
            int found = devgroup_member(inst->topo_sqlite_db, node, id, name);
            if (found < 0) return identity_error(inst, "subgroup membership verification failed", node);
            if (found) {
                struct devgroup_record record; struct devgroup_profile profile;
                if (devgroup_latest(inst->topo_sqlite_db, id, &record, &profile) != 1)
                    return identity_error(inst, "subgroup profile unavailable", node);
                party.kind = DM_PARTY_GROUP; party.node = 0;
                memcpy(party.ed, record.ed, 32); memcpy(party.x, record.x, 32);
            }
        }
    }
    if (!party_valid(&party)) return identity_error(inst, "invalid resolved party", node);
    *out = party;
    DEBUG_DEBUG(DEBUG_CATEGORY_DM, "DM identity: resolved node=%llu mode=%d kind=%u",
        (unsigned long long)node, mode, party.kind);
    return 0;
}

int dm_party_member(struct UTUN_INSTANCE* inst, const struct dm_party* party, uint64_t node) {
    if (!party_valid(party) || !node) return identity_error(inst, "invalid membership arguments", node);
    if (party->kind == DM_PARTY_DEVICE && party->node != node) return 0;
    struct dm_party current;
    if (dm_party_resolve(inst, node, party->kind == DM_PARTY_GROUP ? DM_MODE_GROUPS : DM_MODE_DEVICE, &current)) return -1;
    return dm_party_equal(party, &current);
}

int dm_party_private(struct UTUN_INSTANCE* inst, const struct dm_party* party, uint8_t out[32]) {
    if (!inst || !out || !party_valid(party)) return identity_error(inst, "invalid secret request", 0);
    if (party->kind == DM_PARTY_DEVICE) {
        if (party->node != inst->node_id || memcmp(party->x, inst->my_keys.public_key, 32) ||
            memcmp(party->ed, inst->my_ed25519_pubkey, 32)) return identity_error(inst, "device secret unavailable", party->node);
        memcpy(out, inst->my_keys.private_key, 32);
        return 0;
    }
    struct devgroup_record record; uint8_t ed[32] = {0}, x[32] = {0};
    int rc = devgroup_load_local(inst, &record, ed, x);
    int found = rc == 1 && !memcmp(party->ed, record.ed, 32) && !memcmp(party->x, record.x, 32);
    if (found) memcpy(out, x, 32);
    OPENSSL_cleanse(ed, sizeof(ed)); OPENSSL_cleanse(x, sizeof(x));
    if (!found) return identity_error(inst, "pinned subgroup secret unavailable", inst->node_id);
    return 0;
}

uint64_t dm_identity_conv_id(enum dm_mode mode, const struct dm_party parties[2]) {
    if (!parties || (mode != DM_MODE_DEVICE && mode != DM_MODE_GROUPS) || !party_valid(parties) ||
        !party_valid(parties + 1) || dm_party_equal(parties, parties + 1)) {
        identity_error(NULL, "invalid conversation parties", 0); return 0;
    }
    if (mode == DM_MODE_DEVICE) {
        if (parties[0].kind != DM_PARTY_DEVICE || parties[1].kind != DM_PARTY_DEVICE || parties[0].node == parties[1].node) {
            identity_error(NULL, "device mode requires distinct devices", 0); return 0;
        }
        return dm_derive_conv_id(parties[0].node, parties[1].node);
    }
    const char domain[] = "utun-dm-groups-v1";
    uint8_t wire[2][DM_PARTY_WIRE_SIZE], data[sizeof(domain) + 2 * DM_PARTY_WIRE_SIZE], hash[32];
    if (dm_party_encode(parties, wire[0]) || dm_party_encode(parties + 1, wire[1])) return 0;
    unsigned first = memcmp(wire[0], wire[1], DM_PARTY_WIRE_SIZE) > 0;
    memcpy(data, domain, sizeof(domain)); memcpy(data + sizeof(domain), wire[first], DM_PARTY_WIRE_SIZE);
    memcpy(data + sizeof(domain) + DM_PARTY_WIRE_SIZE, wire[1 - first], DM_PARTY_WIRE_SIZE);
    SHA256(data, sizeof(data), hash);
    uint64_t id = 0;
    for (unsigned i = 0; i < 8; i++) id |= (uint64_t)hash[i] << (8 * i);
    id &= INT64_MAX;
    return id ? id : 1;
}

/* Добавление не молча обрезает snapshot и не допускает принадлежность обеим сторонам. */
static int targets_add(struct dm_target targets[DM_TARGET_MAX], size_t* count, size_t capacity, uint64_t node, unsigned side) {
    for (size_t i = 0; i < *count; i++) {
        if (targets[i].node != node) continue;
        if (targets[i].side == side) return 0;
        return identity_error(NULL, "device belongs to both parties", node);
    }
    if (*count >= capacity || *count >= DM_TARGET_MAX) return identity_error(NULL, "recipient snapshot limit exceeded", node);
    targets[*count].node = node; targets[*count].side = (uint8_t)side; ++*count;
    return 0;
}

static int target_compare(const void* a, const void* b) {
    uint64_t x = ((const struct dm_target*)a)->node, y = ((const struct dm_target*)b)->node;
    return x > y ? 1 : x < y ? -1 : 0;
}

int dm_identity_targets(struct UTUN_INSTANCE* inst, const struct dm_party parties[2],
    struct dm_target* out, size_t capacity, size_t* count) {
    if (!inst || !inst->topo_sqlite_db || !parties || !out || !count || !capacity ||
        !party_valid(parties) || !party_valid(parties + 1) || dm_party_equal(parties, parties + 1))
        return identity_error(inst, "invalid recipient snapshot arguments", 0);
    struct dm_target targets[DM_TARGET_MAX]; size_t n = 0;
    for (unsigned side = 0; side < 2; side++) {
        const struct dm_party* party = parties + side;
        if (party->kind == DM_PARTY_DEVICE) {
            if (dm_party_member(inst, party, party->node) != 1 || targets_add(targets, &n, capacity, party->node, side)) return -1;
            continue;
        }
        size_t before = n;
        int local = dm_party_member(inst, party, inst->node_id);
        if (local < 0 || (local && targets_add(targets, &n, capacity, inst->node_id, side))) return -1;
        sqlite3_stmt* st = NULL;
        int rc = sqlite3_prepare_v2(inst->topo_sqlite_db, "SELECT node_id FROM nodes WHERE devgroup_ed=?", -1, &st, NULL);
        if (rc == SQLITE_OK) rc = sqlite3_bind_blob(st, 1, party->ed, 32, SQLITE_STATIC);
        if (rc == SQLITE_OK) while ((rc = sqlite3_step(st)) == SQLITE_ROW) {
            uint64_t node = (uint64_t)sqlite3_column_int64(st, 0);
            if (node == inst->node_id) continue;
            int member = dm_party_member(inst, party, node);
            if (member != 1 || targets_add(targets, &n, capacity, node, side)) {
                sqlite3_finalize(st); return identity_error(inst, "inconsistent recipient membership", node);
            }
        }
        int final = sqlite3_finalize(st);
        if (rc != SQLITE_DONE || final != SQLITE_OK) return identity_error(inst, "recipient query failed", 0);
        if (n == before) return identity_error(inst, "pinned party has no known devices", 0);
    }
    qsort(targets, n, sizeof(targets[0]), target_compare);
    memcpy(out, targets, n * sizeof(targets[0])); *count = n;
    DEBUG_DEBUG(DEBUG_CATEGORY_DM, "DM identity: recipient snapshot complete count=%zu", n);
    return 0;
}