#include "../../protocol_core/BLEPeerSessionManager.h"

#include <cassert>
#include <cstdint>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>

using namespace ble_reticulum;

namespace {

PeerIdentity identity_with_base(uint8_t base) {
    PeerIdentity identity{};
    for (size_t i = 0; i < identity.size(); ++i) {
        identity[i] = static_cast<uint8_t>(base + i);
    }
    return identity;
}

std::vector<uint8_t> to_payload(const PeerIdentity& identity) {
    return std::vector<uint8_t>(identity.begin(), identity.end());
}

ConnectionSnapshot snapshot(const std::string& address) {
    ConnectionSnapshot snap;
    snap.current.address = address;
    snap.local_role = LocalRole::Peripheral;
    return snap;
}

bool has_action(const HandshakeResult& result, SessionActionType type) {
    for (const SessionAction& action : result.actions) {
        if (action.type == type) {
            return true;
        }
    }
    return false;
}

void test_non_16_byte_payload() {
    BLEPeerSessionManager manager;
    auto snap = snapshot("AA:BB:CC:00:00:01");
    std::vector<uint8_t> data{1, 2, 3};

    HandshakeResult result =
        manager.handleIdentityHandshake(snap, data.data(), data.size(), 10.0);

    assert(result.decision == InputDecision::PassToReassembler);
    assert(!result.consumed);
    assert(!result.accepted);
    assert(has_action(result, SessionActionType::PassToReassembler));
}

void test_new_16_byte_identity() {
    BLEPeerSessionManager manager;
    PeerIdentity identity = identity_with_base(0x10);
    std::vector<uint8_t> data = to_payload(identity);
    auto snap = snapshot("AA:BB:CC:00:00:02");

    HandshakeResult result =
        manager.handleIdentityHandshake(snap, data.data(), data.size(), 11.0);

    assert(result.decision == InputDecision::AcceptedNewIdentity);
    assert(result.consumed);
    assert(result.accepted);
    assert(result.identity_key == "1011121314151617");
    assert(result.fragmenter_key == "101112131415161718191a1b1c1d1e1f");
    assert(has_action(result, SessionActionType::AcceptNewIdentity));
    assert(has_action(result, SessionActionType::CreateFragmentationState));
    assert(has_action(result, SessionActionType::MarkPeerReady));
}

void test_known_identity_duplicate_same() {
    BLEPeerSessionManager manager;
    PeerIdentity identity = identity_with_base(0x20);
    std::vector<uint8_t> data = to_payload(identity);
    auto snap = snapshot("AA:BB:CC:00:00:03");
    snap.known_identity_for_address = identity;

    HandshakeResult result =
        manager.handleIdentityHandshake(snap, data.data(), data.size(), 12.0);

    assert(result.decision == InputDecision::ConsumedDuplicateSameIdentity);
    assert(result.consumed);
    assert(!result.accepted);
}

void test_known_identity_duplicate_mismatch() {
    BLEPeerSessionManager manager;
    PeerIdentity known = identity_with_base(0x30);
    PeerIdentity incoming = identity_with_base(0x40);
    std::vector<uint8_t> data = to_payload(incoming);
    auto snap = snapshot("AA:BB:CC:00:00:04");
    snap.known_identity_for_address = known;

    HandshakeResult result =
        manager.handleIdentityHandshake(snap, data.data(), data.size(), 13.0);

    assert(result.decision == InputDecision::ConsumedDuplicateMismatchedIdentity);
    assert(result.consumed);
    assert(has_action(result, SessionActionType::Warn));
}

void test_duplicate_identity_active_elsewhere() {
    BLEPeerSessionManager manager;
    PeerIdentity identity = identity_with_base(0x50);
    std::vector<uint8_t> data = to_payload(identity);
    auto first = snapshot("AA:BB:CC:00:00:05");
    manager.handleIdentityHandshake(first, data.data(), data.size(), 14.0);

    auto duplicate = snapshot("AA:BB:CC:00:00:06");
    duplicate.existing_address_for_identity = first.current.address;
    duplicate.existing_address_connected = true;
    duplicate.existing_address_in_peer_table = true;

    HandshakeResult result =
        manager.handleIdentityHandshake(duplicate, data.data(), data.size(), 15.0);

    assert(result.decision == InputDecision::RejectedDuplicateIdentity);
    assert(result.consumed);
    assert(!result.accepted);
    assert(result.should_disconnect_current);
    assert(has_action(result, SessionActionType::DisconnectCurrentPeer));
}

void test_duplicate_identity_with_stale_pending_detach() {
    BLEPeerSessionManager manager;
    PeerIdentity identity = identity_with_base(0x60);
    std::vector<uint8_t> data = to_payload(identity);
    auto first = snapshot("AA:BB:CC:00:00:07");
    manager.handleIdentityHandshake(first, data.data(), data.size(), 16.0);

    auto rotated = snapshot("AA:BB:CC:00:00:08");
    rotated.existing_address_for_identity = first.current.address;
    rotated.identity_has_pending_detach = true;

    HandshakeResult result =
        manager.handleIdentityHandshake(rotated, data.data(), data.size(), 17.0);

    assert(result.decision == InputDecision::AcceptedNewIdentity);
    assert(result.accepted);
    assert(has_action(result, SessionActionType::CleanupOldAddress));
    assert(has_action(result, SessionActionType::UpdatePeerAddress));

    auto view = manager.sessionByIdentity(identity);
    assert(view.has_value());
    assert(view->current_address == rotated.current.address);
}

void test_duplicate_identity_with_zombie_old_connection() {
    BLEPeerSessionManager manager;
    PeerIdentity identity = identity_with_base(0x70);
    std::vector<uint8_t> data = to_payload(identity);
    auto first = snapshot("AA:BB:CC:00:00:09");
    manager.handleIdentityHandshake(first, data.data(), data.size(), 18.0);

    auto replacement = snapshot("AA:BB:CC:00:00:0A");
    replacement.existing_address_for_identity = first.current.address;
    replacement.existing_address_connected = true;
    replacement.existing_address_in_peer_table = true;
    replacement.existing_connection_is_zombie = true;

    HandshakeResult result =
        manager.handleIdentityHandshake(replacement, data.data(), data.size(), 19.0);

    assert(result.decision == InputDecision::AcceptedNewIdentity);
    assert(result.accepted);
    assert(result.should_disconnect_old);
    assert(has_action(result, SessionActionType::DisconnectOldPeer));
    assert(has_action(result, SessionActionType::CleanupOldAddress));
}

void test_mtu_provided() {
    BLEPeerSessionManager manager;
    PeerIdentity identity = identity_with_base(0x80);
    std::vector<uint8_t> data = to_payload(identity);
    auto snap = snapshot("AA:BB:CC:00:00:0B");
    snap.negotiated_mtu = 185;

    HandshakeResult result =
        manager.handleIdentityHandshake(snap, data.data(), data.size(), 20.0);

    assert(result.mtu == 185);
    auto view = manager.sessionByIdentity(identity);
    assert(view.has_value());
    assert(view->mtu == 185);
}

void test_mtu_missing_fallback() {
    BLEPeerSessionManager manager;
    PeerIdentity identity = identity_with_base(0x90);
    std::vector<uint8_t> data = to_payload(identity);
    auto snap = snapshot("AA:BB:CC:00:00:0C");

    HandshakeResult result =
        manager.handleIdentityHandshake(snap, data.data(), data.size(), 21.0);

    assert(result.mtu == 23);
}

void test_pending_identity_timeout() {
    BLEPeerSessionManager manager(30.0, 45.0);
    ConnectionId one{"AA:BB:CC:00:00:0D", 1};
    ConnectionId two{"AA:BB:CC:00:00:0E", 2};
    manager.markPendingIdentity(one, 100.0);
    manager.markPendingIdentity(two, 120.0);

    std::vector<ConnectionId> expired = manager.expiredPendingIdentities(131.0);
    assert(expired.size() == 1);
    assert(expired[0] == one);

    manager.removePendingIdentity(one);
    expired = manager.expiredPendingIdentities(200.0);
    assert(expired.size() == 1);
    assert(expired[0] == two);
}

void test_peer_address_update_fragmenter_key_unchanged() {
    BLEPeerSessionManager manager;
    PeerIdentity identity = identity_with_base(0xa0);
    std::vector<uint8_t> data = to_payload(identity);
    auto first = snapshot("AA:BB:CC:00:00:0F");
    HandshakeResult first_result =
        manager.handleIdentityHandshake(first, data.data(), data.size(), 22.0);

    auto second = snapshot("AA:BB:CC:00:00:10");
    second.existing_address_for_identity = first.current.address;
    second.identity_has_pending_detach = true;
    HandshakeResult second_result =
        manager.handleIdentityHandshake(second, data.data(), data.size(), 23.0);

    assert(first_result.fragmenter_key == second_result.fragmenter_key);
    auto old_view = manager.sessionByAddress(first.current.address);
    auto new_view = manager.sessionByAddress(second.current.address);
    assert(!old_view.has_value());
    assert(new_view.has_value());
    assert(new_view->current_address == second.current.address);
}

void test_key_helpers() {
    PeerIdentity identity = identity_with_base(0x01);
    assert(BLEPeerSessionManager::computeIdentityKey(identity) == "0102030405060708");
    assert(BLEPeerSessionManager::computeFragmenterKey(identity) ==
           "0102030405060708090a0b0c0d0e0f10");
}

void test_identity_from_payload_rejects_non_16() {
    std::vector<uint8_t> data{1, 2, 3};
    bool threw = false;
    try {
        (void)BLEPeerSessionManager::identityFromPayload(data.data(), data.size());
    } catch (const std::invalid_argument&) {
        threw = true;
    }
    assert(threw);
}

} // namespace

int main() {
    test_non_16_byte_payload();
    test_new_16_byte_identity();
    test_known_identity_duplicate_same();
    test_known_identity_duplicate_mismatch();
    test_duplicate_identity_active_elsewhere();
    test_duplicate_identity_with_stale_pending_detach();
    test_duplicate_identity_with_zombie_old_connection();
    test_mtu_provided();
    test_mtu_missing_fallback();
    test_pending_identity_timeout();
    test_peer_address_update_fragmenter_key_unchanged();
    test_key_helpers();
    test_identity_from_payload_rejects_non_16();

    std::cout << "BLEPeerSessionManager native tests passed\n";
    return 0;
}