#include "BLEPeerSessionManager.h"

#include <algorithm>
#include <stdexcept>

namespace ble_reticulum {

namespace {

constexpr uint16_t BLE_MINIMUM_MTU = 23;

SessionAction make_action(SessionActionType type, const ConnectionId& target) {
    SessionAction action;
    action.type = type;
    action.target = target;
    return action;
}

SessionAction make_message_action(SessionActionType type,
                                  const ConnectionId& target,
                                  const std::string& message) {
    SessionAction action = make_action(type, target);
    action.message = message;
    return action;
}

} // namespace

BLEPeerSessionManager::BLEPeerSessionManager(double pending_identity_timeout,
                                             double zombie_timeout)
    : pending_identity_timeout_(pending_identity_timeout),
      zombie_timeout_(zombie_timeout) {}

bool BLEPeerSessionManager::isIdentityHandshakePayload(const uint8_t* data,
                                                       size_t data_size) {
    return data != nullptr && data_size == 16;
}

PeerIdentity BLEPeerSessionManager::identityFromPayload(const uint8_t* data,
                                                        size_t data_size) {
    if (!isIdentityHandshakePayload(data, data_size)) {
        throw std::invalid_argument("identity payload must be exactly 16 bytes");
    }

    PeerIdentity identity{};
    std::copy(data, data + identity.size(), identity.begin());
    return identity;
}

std::string BLEPeerSessionManager::computeIdentityKey(const PeerIdentity& identity) {
    static constexpr char hex[] = "0123456789abcdef";
    std::string out;
    out.reserve(16);
    for (size_t i = 0; i < 8; ++i) {
        const uint8_t byte = identity[i];
        out.push_back(hex[(byte >> 4) & 0x0f]);
        out.push_back(hex[byte & 0x0f]);
    }
    return out;
}

std::string BLEPeerSessionManager::computeFragmenterKey(const PeerIdentity& identity) {
    static constexpr char hex[] = "0123456789abcdef";
    std::string out;
    out.reserve(32);
    for (uint8_t byte : identity) {
        out.push_back(hex[(byte >> 4) & 0x0f]);
        out.push_back(hex[byte & 0x0f]);
    }
    return out;
}

HandshakeResult BLEPeerSessionManager::handleIdentityHandshake(
    const ConnectionSnapshot& connection,
    const uint8_t* data,
    size_t data_size,
    double now_seconds) {

    HandshakeResult result;
    result.mtu = connection.negotiated_mtu.value_or(BLE_MINIMUM_MTU);

    if (!isIdentityHandshakePayload(data, data_size)) {
        result.decision = InputDecision::PassToReassembler;
        result.consumed = false;
        result.actions.push_back(make_action(SessionActionType::PassToReassembler,
                                             connection.current));
        return result;
    }

    PeerIdentity payload_identity = identityFromPayload(data, data_size);
    result.peer_identity = payload_identity;
    result.identity_key = computeIdentityKey(payload_identity);
    result.fragmenter_key = computeFragmenterKey(payload_identity);
    result.consumed = true;

    if (connection.known_identity_for_address.has_value()) {
        if (connection.known_identity_for_address.value() == payload_identity) {
            result.decision = InputDecision::ConsumedDuplicateSameIdentity;
            result.actions.push_back(make_action(SessionActionType::ConsumeInput,
                                                 connection.current));
            return result;
        }

        result.decision = InputDecision::ConsumedDuplicateMismatchedIdentity;
        result.actions.push_back(make_message_action(
            SessionActionType::Warn,
            connection.current,
            "16-byte data differs from known identity; consumed as identity-like data"));
        result.actions.push_back(make_action(SessionActionType::ConsumeInput,
                                             connection.current));
        return result;
    }

    const PeerSession* existing_session = findSessionByIdentityKey(result.identity_key);
    std::optional<std::string> existing_address = connection.existing_address_for_identity;
    if (!existing_address.has_value() && existing_session &&
        !existing_session->current_address.empty()) {
        existing_address = existing_session->current_address;
    }

    const bool duplicate_elsewhere =
        existing_address.has_value() && existing_address.value() != connection.current.address;

    if (duplicate_elsewhere) {
        const bool stale_or_pending =
            connection.identity_has_pending_detach ||
            (!connection.existing_address_connected &&
             !connection.existing_address_in_peer_table);

        if (connection.existing_connection_is_zombie || stale_or_pending) {
            SessionAction cleanup = make_action(SessionActionType::CleanupOldAddress,
                                                connection.current);
            cleanup.old_address = existing_address.value();
            cleanup.new_address = connection.current.address;
            result.actions.push_back(cleanup);

            if (connection.existing_connection_is_zombie) {
                result.should_disconnect_old = true;
                SessionAction disconnect_old =
                    make_action(SessionActionType::DisconnectOldPeer, connection.current);
                disconnect_old.old_address = existing_address.value();
                result.actions.push_back(disconnect_old);
            }

            SessionAction update = make_action(SessionActionType::UpdatePeerAddress,
                                               connection.current);
            update.old_address = existing_address.value();
            update.new_address = connection.current.address;
            result.actions.push_back(update);
        } else {
            result.decision = InputDecision::RejectedDuplicateIdentity;
            result.should_disconnect_current = true;
            result.actions.push_back(make_action(SessionActionType::RejectDuplicateIdentity,
                                                 connection.current));
            result.actions.push_back(make_action(SessionActionType::DisconnectCurrentPeer,
                                                 connection.current));
            return result;
        }
    }

    upsertAcceptedSession(connection, payload_identity, result.mtu, now_seconds);
    removePendingIdentity(connection.current);

    result.decision = InputDecision::AcceptedNewIdentity;
    result.accepted = true;
    result.actions.push_back(make_action(SessionActionType::AcceptNewIdentity,
                                         connection.current));
    result.actions.push_back(make_action(SessionActionType::CreateFragmentationState,
                                         connection.current));
    result.actions.push_back(make_action(SessionActionType::MarkPeerReady,
                                         connection.current));
    result.actions.push_back(make_action(SessionActionType::RemovePendingIdentity,
                                         connection.current));
    result.actions.push_back(make_action(SessionActionType::MarkRealData,
                                         connection.current));
    return result;
}

void BLEPeerSessionManager::markConnected(const ConnectionSnapshot& connection,
                                          double now_seconds) {
    if (connection.known_identity_for_address.has_value()) {
        upsertAcceptedSession(connection,
                              connection.known_identity_for_address.value(),
                              connection.negotiated_mtu.value_or(BLE_MINIMUM_MTU),
                              now_seconds);
    } else {
        markPendingIdentity(connection.current, now_seconds);
    }
}

void BLEPeerSessionManager::markDisconnected(const ConnectionId& connection,
                                             double /*now_seconds*/) {
    sessions_.erase(std::remove_if(sessions_.begin(), sessions_.end(),
                                   [&](const PeerSession& session) {
                                       return session.current_address == connection.address ||
                                              session.current_handle == connection.handle;
                                   }),
                    sessions_.end());
    removePendingIdentity(connection);
}

void BLEPeerSessionManager::markMtu(const ConnectionId& connection, uint16_t mtu) {
    PeerSession* session = findSessionByAddress(connection.address);
    if (session) {
        session->mtu = mtu;
    }
}

void BLEPeerSessionManager::markPendingIdentity(const ConnectionId& connection,
                                                double now_seconds) {
    for (PendingIdentity& pending : pending_identities_) {
        if (pending.connection.address == connection.address) {
            pending.connection = connection;
            pending.started_at = now_seconds;
            return;
        }
    }

    PendingIdentity pending;
    pending.connection = connection;
    pending.started_at = now_seconds;
    pending_identities_.push_back(pending);
}

void BLEPeerSessionManager::removePendingIdentity(const ConnectionId& connection) {
    pending_identities_.erase(
        std::remove_if(pending_identities_.begin(), pending_identities_.end(),
                       [&](const PendingIdentity& pending) {
                           return pending.connection.address == connection.address;
                       }),
        pending_identities_.end());
}

std::vector<ConnectionId> BLEPeerSessionManager::expiredPendingIdentities(
    double now_seconds) const {
    std::vector<ConnectionId> expired;
    for (const PendingIdentity& pending : pending_identities_) {
        if (now_seconds - pending.started_at > pending_identity_timeout_) {
            expired.push_back(pending.connection);
        }
    }
    return expired;
}

std::optional<PeerSessionView> BLEPeerSessionManager::sessionByAddress(
    const std::string& address) const {
    const PeerSession* session = findSessionByAddress(address);
    if (!session) {
        return std::nullopt;
    }
    PeerSessionView view;
    view.identity = session->identity;
    view.identity_key = session->identity_key;
    view.current_address = session->current_address;
    view.current_handle = session->current_handle;
    view.mtu = session->mtu;
    view.has_fragmentation_state = session->has_fragmentation_state;
    view.peer_ready = session->peer_ready;
    view.pending_identity_since = session->pending_identity_since;
    view.last_real_data = session->last_real_data;
    return view;
}

std::optional<PeerSessionView> BLEPeerSessionManager::sessionByIdentity(
    const PeerIdentity& identity) const {
    const PeerSession* session = findSessionByIdentityKey(computeIdentityKey(identity));
    if (!session) {
        return std::nullopt;
    }
    PeerSessionView view;
    view.identity = session->identity;
    view.identity_key = session->identity_key;
    view.current_address = session->current_address;
    view.current_handle = session->current_handle;
    view.mtu = session->mtu;
    view.has_fragmentation_state = session->has_fragmentation_state;
    view.peer_ready = session->peer_ready;
    view.pending_identity_since = session->pending_identity_since;
    view.last_real_data = session->last_real_data;
    return view;
}

BLEPeerSessionManager::PeerSession* BLEPeerSessionManager::findSessionByIdentityKey(
    const std::string& identity_key) {
    for (PeerSession& session : sessions_) {
        if (session.identity_key == identity_key) {
            return &session;
        }
    }
    return nullptr;
}

const BLEPeerSessionManager::PeerSession* BLEPeerSessionManager::findSessionByIdentityKey(
    const std::string& identity_key) const {
    for (const PeerSession& session : sessions_) {
        if (session.identity_key == identity_key) {
            return &session;
        }
    }
    return nullptr;
}

BLEPeerSessionManager::PeerSession* BLEPeerSessionManager::findSessionByAddress(
    const std::string& address) {
    for (PeerSession& session : sessions_) {
        if (session.current_address == address) {
            return &session;
        }
    }
    return nullptr;
}

const BLEPeerSessionManager::PeerSession* BLEPeerSessionManager::findSessionByAddress(
    const std::string& address) const {
    for (const PeerSession& session : sessions_) {
        if (session.current_address == address) {
            return &session;
        }
    }
    return nullptr;
}

void BLEPeerSessionManager::upsertAcceptedSession(
    const ConnectionSnapshot& connection,
    const PeerIdentity& identity,
    uint16_t mtu,
    double now_seconds) {
    const std::string identity_key = computeIdentityKey(identity);
    PeerSession* session = findSessionByIdentityKey(identity_key);
    if (!session) {
        PeerSession new_session;
        new_session.identity = identity;
        new_session.identity_key = identity_key;
        sessions_.push_back(new_session);
        session = &sessions_.back();
    }

    session->identity = identity;
    session->identity_key = identity_key;
    session->current_address = connection.current.address;
    session->current_handle = connection.current.handle;
    session->mtu = mtu;
    session->has_fragmentation_state = true;
    session->peer_ready = true;
    session->pending_identity_since = 0.0;
    session->last_real_data = now_seconds;
}

} // namespace ble_reticulum