/*
 *  Copyright (C) 2004-2026 Savoir-faire Linux Inc.
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program. If not, see <https://www.gnu.org/licenses/>.
 */
#pragma once

#include "routing_table.h"
#include "swarm_protocol.h"

#include <memory>

namespace jami {

using namespace swarm_protocol;

class SwarmManager : public std::enable_shared_from_this<SwarmManager>
{
    using ChannelCb = std::function<bool(const std::shared_ptr<dhtnet::ChannelSocketInterface>&)>;
    using NeedSocketCb = std::function<void(const std::string&, ChannelCb&&, bool noNewSocket)>;
    using ToConnectCb = std::function<bool(const NodeId&)>;
    using OnConnectionChanged = std::function<void(bool ok)>;
    using OnMobileNodesChanged = std::function<void(const std::vector<NodeId>&)>;
    using OnMobileNodeInfosChanged = std::function<void(const std::vector<MobileNodeInfo>&)>;
    using MobileLeaseProvider = std::function<std::optional<MobileNodeInfo>()>;
    using MobileLeaseIssuerValidator = std::function<bool(const dht::InfoHash&)>;
    /** Synchronous lookup of an already known device certificate (account certificate store). */
    using CertificateProvider = std::function<std::shared_ptr<dht::crypto::Certificate>(const NodeId&)>;
    /** Asynchronous last-resort lookup (DHT). The implementation is expected to pin on success. */
    using CertificateFetcher
        = std::function<void(const NodeId&, std::function<void(const std::shared_ptr<dht::crypto::Certificate>&)>&&)>;

public:
    explicit SwarmManager(const NodeId& nodeId,
                          bool isMobile,
                          const std::mt19937_64& rand,
                          ToConnectCb&& toConnectCb,
                          std::string conversationId = {},
                          MobileLeaseProvider mobileLeaseProvider = {},
                          MobileLeaseIssuerValidator mobileLeaseIssuerValidator = {},
                          CertificateProvider certificateProvider = {},
                          CertificateFetcher certificateFetcher = {});
    ~SwarmManager();

    NeedSocketCb needSocketCb_;

    std::weak_ptr<SwarmManager> weak() { return weak_from_this(); }

    /**
     * Get swarm manager id
     * @return NodeId
     */
    const NodeId& getId() const { return id_; }

    /**
     * Set list of nodes to the routing table known_nodes
     * @param known_nodes
     * @return if some are new
     */
    bool setKnownNodes(const std::vector<NodeId>& known_nodes);

    /**
     * Set list of nodes to the routing table mobile_nodes
     * @param mobile_nodes
     */
    void setMobileNodes(const std::vector<NodeId>& mobile_nodes);

    /**
     * Merge lease-bearing mobile node records received from peers.
     * Records whose lease cannot be verified yet (the device certificate is
     * unknown locally) are kept aside and resolved asynchronously; they only
     * become visible once verified.
     * @param mobile_nodes
     */
    void setMobileNodes(const std::vector<MobileNodeInfo>& mobile_nodes, bool requireLease = false);

    /**
     * Add channel to routing table
     * @param channel
     */
    void addChannel(const std::shared_ptr<dhtnet::ChannelSocketInterface>& channel);

    /**
     * Remove channel from routing table
     * @param channel
     */
    void removeNode(const NodeId& nodeId);

    /**
     * Change mobility of specific node
     * @param nodeId
     * @param isMobile
     */
    void changeMobility(const NodeId& nodeId, bool isMobile);

    /**
     * Check if swarm manager is connected
     * @return true if the swarm has at least one connected node, false if not
     */
    bool isConnected() const;

    /**
     * get all nodes from the different tables in bucket
     */
    std::vector<NodeId> getAllNodes() const;

    std::vector<NodeId> getConnectedNodes() const;

    /**
     * Get the mobile nodes this device is responsible for waking up.
     * The closest redundant nodes in the local Kademlia view wake it.
     */
    std::vector<NodeId> getMobileNodesToNotify();

    /**
     * Get every node known to be mobile (connected or not).
     * Used to persist mobility knowledge across restarts.
     */
    std::vector<NodeId> getKnownMobileNodes() const;

    /**
     * Get known mobile nodes with any validated lease metadata.
     */
    std::vector<MobileNodeInfo> getKnownMobileNodeInfos() const;

    /**
     * Get wake-up targets with any validated lease metadata.
     */
    std::vector<MobileNodeInfo> getMobileNodeInfosToNotify();

    /**
     * Callback invoked when the set of known mobile nodes changes,
     * with the updated set.
     * @param cb
     */
    void onMobileNodesChanged(OnMobileNodesChanged cb)
    {
        std::lock_guard lock(onMobileNodesChangedMtx_);
        onMobileNodesChanged_ = std::move(cb);
    }

    void onMobileNodeInfosChanged(OnMobileNodeInfosChanged cb)
    {
        std::lock_guard lock(onMobileNodesChangedMtx_);
        onMobileNodeInfosChanged_ = std::move(cb);
    }

    std::vector<std::map<std::string, std::string>> getRoutingTableInfo() const;

    unsigned getActiveNodesCount() const { return routing_table.getActiveNodesCount(); }

    /**
     * Delete nodes from the different tables in bucket
     */
    void deleteNode(const std::vector<NodeId>& nodes);

    // For tests

    /**
     * Get routing table
     * @return RoutingTable
     */
    RoutingTable& getRoutingTable() { return routing_table; };

    std::optional<MobileNodeInfo> getLocalMobileNodeInfo() { return localMobileNodeInfo(); }

    /**
     * Get buckets of routing table
     * @return buckets list
     */
    std::list<Bucket>& getBuckets() { return routing_table.getBuckets(); };

    /**
     * Shutdown swarm manager
     */
    void shutdown();

    /**
     * Restart the swarm manager.
     *
     * This function must be called in situations where we want
     * to use a swarm manager that was previously shut down.
     */
    void restart();

    /**
     * Display swarm manager info
     */
    void display()
    {
        JAMI_DEBUG("SwarmManager {:s} has {:d} nodes in table [P = {}]",
                   getId().to_c_str(),
                   routing_table.getNodeCount(),
                   isMobile_);
        // print nodes of routingtable
        for (auto& bucket : routing_table.getBuckets()) {
            for (auto& node : bucket.getNodes()) {
                JAMI_DEBUG("Node {:s}", node.first.toString());
            }
        }
    }

    /*
     * Callback for connection changed
     * @param cb
     */
    void onConnectionChanged(OnConnectionChanged cb) { onConnectionChanged_ = std::move(cb); }

    /**
     * Get mobility of swarm manager
     * @return true if mobile, false if not
     */
    bool isMobile() const { return isMobile_; }

    /**
     * Maintain/Update buckets
     * @param toConnect         Nodes to connect
     */
    void maintainBuckets(const std::set<NodeId>& toConnect = {});

    /**
     * Proactively connect to a node, bypassing bucket capacity checks.
     * The node is registered as known and a connection is attempted with
     * noNewSocket=true (reuses an existing transport).
     * @param nodeId
     */
    void connectNode(const NodeId& nodeId);

    /**
     * Check if swarm manager is shutdown
     * @return true if shutdown, false if not
     */
    bool isShutdown() { return isShutdown_; };

private:
    /**
     * Check if we're connected with a specific device
     * @param deviceId
     * @return true if connected, false if not
     */
    bool isConnectedWith(const NodeId& deviceId);

    /**
     * Add node to the known_nodes list
     * @param nodeId
     * @return if node inserted
     */
    bool addKnownNode(const NodeId& nodeId);

    /**
     * Add node to the mobile_Nodes list
     * @param nodeId
     * @return if node inserted
     */
    bool addMobileNodes(const NodeId& nodeId);

    bool setMobileNodeInfo(const MobileNodeInfo& mobile,
                           bool requireLease = false,
                           const std::shared_ptr<dhtnet::ChannelSocketInterface>& source = {});

    /**
     * Certificate-free validation of a lease. Cheap: no cryptography, no I/O.
     * Gates every expensive certificate resolution, in particular by requiring
     * the issuer to be a member of the conversation.
     */
    bool precheckLease(const MobileLease& lease) const;

    /**
     * Full cryptographic validation of a lease against its device certificate.
     */
    bool verifyLease(const dht::crypto::Certificate& certificate, const MobileLease& lease) const;

    /**
     * Record a verified lease. Must be called with mutex held.
     * @return true if the known mobile node set or leases changed
     */
    bool commitLeaseInternal(const MobileLease& lease);

    /**
     * Keep an unverified lease aside and start resolving its certificate.
     * Must be called with mutex held.
     */
    void enqueuePendingLeaseInternal(const MobileLease& lease,
                                     const std::shared_ptr<dhtnet::ChannelSocketInterface>& source);

    /** Ask a peer for device certificates over the swarm channel. */
    void requestCertificates(const std::shared_ptr<dhtnet::ChannelSocketInterface>& socket,
                             const std::vector<NodeId>& ids);

    void onCertRequest(const std::shared_ptr<dhtnet::ChannelSocketInterface>& socket, const CertRequest& request);

    void onCertResponse(const std::shared_ptr<dhtnet::ChannelSocketInterface>& socket, const CertResponse& response);

    /** Try to promote pending leases now that `certificate` is known. */
    void onCertificateResolved(const NodeId& nodeId, const std::shared_ptr<dht::crypto::Certificate>& certificate);

    /** Last-resort DHT lookup for a pending lease's certificate. */
    void fetchCertificateFromDht(const NodeId& nodeId);

    /**
     * Forget an unverifiable lease: the node is simply ignored until it is
     * announced again. Must be called with mutex held.
     */
    void abandonLeaseInternal(const NodeId& nodeId);

    bool isMobileNodeCurrentInternal(const NodeId& nodeId, int64_t now) const;

    std::optional<MobileNodeInfo> localMobileNodeInfo();

    void scheduleMobileLeaseExpiryInternal();

    void expireMobileLeases(const asio::error_code& ec);

    /**
     * Notify the onMobileNodesChanged_ callback with the current set of
     * known mobile nodes. Must be called without the mutex held.
     */
    void emitMobileNodesChanged();

    /**
     * Send nodes request to fill known_nodes list
     * @param socket
     * @param nodeId
     * @param q
     * @param numberNodes
     */
    void sendRequest(const std::shared_ptr<dhtnet::ChannelSocketInterface>& socket,
                     const NodeId& nodeId,
                     Query q,
                     int numberNodes = Bucket::BUCKET_MAX_SIZE);

    /**
     * Send answer to request
     * @param socket
     * @param msg
     */
    void sendAnswer(const std::shared_ptr<dhtnet::ChannelSocketInterface>& socket, const Message& msg_);

    /**
     * Interpret received message
     * @param socket
     */
    void receiveMessage(const std::shared_ptr<dhtnet::ChannelSocketInterface>& socket);

    /**
     * Reset node's timer expiry
     * @param ec
     * @param socket
     * @param node
     */
    void resetNodeExpiry(const asio::error_code& ec,
                         const std::shared_ptr<dhtnet::ChannelSocketInterface>& socket,
                         NodeId node = {});

    /**
     * Try to establish connection with specific node
     * @param nodeId
     * @param noNewSocket  If true, reuse an existing transport (no new ICE negotiation)
     */
    void tryConnect(const NodeId& nodeId, bool noNewSocket = false);

    /**
     * Remove node from routing table
     * @param nodeId
     */
    void removeNodeInternal(const NodeId& nodeId);

    const NodeId id_;
    bool isMobile_ {false};
    const std::string conversationId_;
    std::mt19937_64 rd;
    mutable std::mutex mutex;
    RoutingTable routing_table;
    std::map<NodeId, MobileLease> mobileNodeLeases_;
    std::map<NodeId, int64_t> legacyMobileNodeExpiries_;

    /** Lease accepted by precheckLease() but not yet verified: certificate unknown. */
    struct PendingLease
    {
        MobileLease lease;
        std::weak_ptr<dhtnet::ChannelSocketInterface> source;
    };
    std::map<NodeId, PendingLease> pendingMobileLeases_;

    /** Nodes whose certificate resolution is already under way: avoids duplicate lookups. */
    std::set<NodeId> certFetchInFlight_;

    /** At most one outstanding CERT_REQ per peer, keyed by the peer's device id. */
    struct CertRequestState
    {
        std::set<NodeId> ids;
        std::shared_ptr<asio::steady_timer> timer;
    };
    std::map<NodeId, CertRequestState> outstandingCertRequests_;

    MobileLeaseProvider mobileLeaseProvider_;
    MobileLeaseIssuerValidator mobileLeaseIssuerValidator_;
    CertificateProvider certificateProvider_;
    CertificateFetcher certificateFetcher_;
    std::mutex mobileLeaseRenewalMtx_;
    std::optional<MobileNodeInfo> localMobileNodeInfo_;
    asio::steady_timer mobileLeaseExpiryTimer_ {*Manager::instance().ioContext()};

    std::atomic_bool isShutdown_ {false};

    OnConnectionChanged onConnectionChanged_ {};
    mutable std::mutex onMobileNodesChangedMtx_;
    std::recursive_mutex mobileNodesEmissionMtx_;
    OnMobileNodesChanged onMobileNodesChanged_ {};
    OnMobileNodeInfosChanged onMobileNodeInfosChanged_ {};

    ToConnectCb toConnectCb_;
};

} // namespace jami
