/*
 *  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/>.
 */
#include "jamidht/message_channel_handler.h"

#include <dhtnet/channel_utils.h>
#include <string_view>

using namespace std::literals;

static constexpr auto MESSAGE_SCHEME = "msg:"sv;

namespace jami {

using Key = std::pair<std::string, DeviceId>;

struct MessageChannelHandler::Impl : public std::enable_shared_from_this<Impl>
{
    dhtnet::ConnectionManager& connectionManager_;
    OnMessage onMessage_;
    OnPeerStateChanged onPeerStateChanged_;
    std::recursive_mutex connectionsMtx_;
    std::map<std::string, std::map<DeviceId, std::vector<std::shared_ptr<dhtnet::ChannelSocket>>>> connections_;

    Impl(dhtnet::ConnectionManager& cm, OnMessage onMessage, OnPeerStateChanged onPeer)
        : connectionManager_(cm)
        , onMessage_(std::move(onMessage))
        , onPeerStateChanged_(std::move(onPeer))
    {}

    void onChannelShutdown(const std::shared_ptr<dhtnet::ChannelSocket>& socket,
                           const std::string& peerId,
                           const DeviceId& device);
};

MessageChannelHandler::MessageChannelHandler(dhtnet::ConnectionManager& cm,
                                             OnMessage onMessage,
                                             OnPeerStateChanged onPeer)
    : ChannelHandlerInterface()
    , pimpl_(std::make_shared<Impl>(cm, std::move(onMessage), std::move(onPeer)))
{}

MessageChannelHandler::~MessageChannelHandler()
{
    std::unique_lock lk(pimpl_->connectionsMtx_);
    for (const auto& [peerId, _] : pimpl_->connections_) {
        pimpl_->onPeerStateChanged_(peerId, false);
    }
    auto connections = std::move(pimpl_->connections_);
    pimpl_->connections_.clear();
    lk.unlock();
}

void
MessageChannelHandler::connect(const DeviceId& deviceId,
                               const std::string&,
                               ConnectCb&& cb,
                               const std::string& connectionType,
                               bool forceNewConnection)
{
    auto channelName = concat(MESSAGE_SCHEME, deviceId.to_view());
    if (pimpl_->connectionManager_.isConnecting(deviceId, channelName)) {
        JAMI_LOG("Already connecting to {}", deviceId);
        return;
    }
    pimpl_->connectionManager_
        .connectDevice(deviceId, channelName, std::move(cb), false, forceNewConnection, connectionType);
}

void
MessageChannelHandler::Impl::onChannelShutdown(const std::shared_ptr<dhtnet::ChannelSocket>& socket,
                                               const std::string& peerId,
                                               const DeviceId& device)
{
    std::lock_guard lk(connectionsMtx_);
    auto peerIt = connections_.find(peerId);
    if (peerIt == connections_.end()) {
        JAMI_WARNING("onChannelShutdown: No connections found for peer {}", peerId);
        return;
    }
    auto connectionsIt = peerIt->second.find(device);
    if (connectionsIt == peerIt->second.end()) {
        JAMI_WARNING("onChannelShutdown: No connections found for device {} of peer {}", device.toString(), peerId);
        return;
    }
    auto& connections = connectionsIt->second;
    auto conn = std::find(connections.begin(), connections.end(), socket);
    if (conn != connections.end())
        connections.erase(conn);
    if (connections.empty()) {
        peerIt->second.erase(connectionsIt);
    }
    if (peerIt->second.empty()) {
        connections_.erase(peerIt);
        onPeerStateChanged_(peerId, false);
    }
}

std::shared_ptr<dhtnet::ChannelSocket>
MessageChannelHandler::getChannel(const std::string& peer, const DeviceId& deviceId) const
{
    std::lock_guard lk(pimpl_->connectionsMtx_);
    auto it = pimpl_->connections_.find(peer);
    if (it == pimpl_->connections_.end())
        return nullptr;
    auto deviceIt = it->second.find(deviceId);
    if (deviceIt == it->second.end())
        return nullptr;
    if (deviceIt->second.empty())
        return nullptr;
    return deviceIt->second.back();
}

std::vector<std::shared_ptr<dhtnet::ChannelSocket>>
MessageChannelHandler::getChannels(const std::string& peer) const
{
    std::vector<std::shared_ptr<dhtnet::ChannelSocket>> sockets;
    std::lock_guard lk(pimpl_->connectionsMtx_);
    auto it = pimpl_->connections_.find(peer);
    if (it == pimpl_->connections_.end())
        return sockets;
    sockets.reserve(it->second.size());
    for (auto& [deviceId, channels] : it->second) {
        for (auto& channel : channels) {
            sockets.push_back(channel);
        }
    }
    return sockets;
}

bool
MessageChannelHandler::onRequest(const std::shared_ptr<dht::crypto::Certificate>& cert, const std::string& /* name */)
{
    if (!cert || !cert->issuer)
        return false;
    return true;
}

void
MessageChannelHandler::onReady(const std::shared_ptr<dht::crypto::Certificate>& cert,
                               const std::string&,
                               std::shared_ptr<dhtnet::ChannelSocket> socket)
{
    if (!cert || !cert->issuer)
        return;
    auto peerId = cert->issuer->getId().toString();
    auto device = cert->getLongId();
    std::lock_guard lk(pimpl_->connectionsMtx_);
    auto& connections = pimpl_->connections_[peerId];
    bool newPeerConnection = connections.empty();
    auto& deviceConnections = connections[device];
    deviceConnections.push_back(socket);
    if (newPeerConnection)
        pimpl_->onPeerStateChanged_(peerId, true);

    socket->setOnRecv(dhtnet::buildMsgpackReader<Message>([onMessage = pimpl_->onMessage_, cert](Message&& msg) {
        onMessage(cert, msg.t, msg.c);
        return std::error_code();
    }));

    socket->onShutdown(
        [w = pimpl_->weak_from_this(), peerId, device, s = std::weak_ptr(socket)](const std::error_code& /*ec*/) {
            if (auto shared = w.lock())
                shared->onChannelShutdown(s.lock(), peerId, device);
        });
}

void
MessageChannelHandler::closeChannel(const std::string& peer,
                                    const DeviceId& device,
                                    const std::shared_ptr<dhtnet::ChannelSocket>& conn)
{
    if (!conn)
        return;
    std::unique_lock lk(pimpl_->connectionsMtx_);
    auto it = pimpl_->connections_.find(peer);
    if (it != pimpl_->connections_.end()) {
        auto deviceIt = it->second.find(device);
        if (deviceIt != it->second.end()) {
            auto& channels = deviceIt->second;
            channels.erase(std::remove(channels.begin(), channels.end(), conn), channels.end());
            if (channels.empty()) {
                it->second.erase(deviceIt);
                if (it->second.empty()) {
                    pimpl_->connections_.erase(it);
                    pimpl_->onPeerStateChanged_(peer, false);
                }
            }
        }
    }
    lk.unlock();
    conn->stop();
}

bool
MessageChannelHandler::sendMessage(const std::shared_ptr<dhtnet::ChannelSocket>& socket, const Message& message)
{
    if (!socket)
        return false;
    msgpack::sbuffer buffer(UINT16_MAX); // Use max
    msgpack::pack(buffer, message);
    std::error_code ec;
    auto sent = socket->write(reinterpret_cast<const uint8_t*>(buffer.data()), buffer.size(), ec);
    if (ec) {
        JAMI_WARNING("Error sending message: {:s}", ec.message());
    }
    return !ec && sent == buffer.size();
}

} // namespace jami
