// tcp_client.cpp - TCP Client for ColumnLynx
// Copyright (C) 2025 DcruBro
// Distributed under the terms of the GNU General Public License, either version 2 only or version 3. See LICENSES/ for details.

#include <columnlynx/client/net/tcp/tcp_client.hpp>
//#include <arpa/inet.h>

namespace ColumnLynx::Net::TCP {
    void TCPClient::start() {
        auto self = shared_from_this();
        mResolver.async_resolve(mHost, mPort,
            [this, self](asio::error_code ec, tcp::resolver::results_type endpoints) {
                if (!ec) {
                    asio::async_connect(mSocket, endpoints,
                        [this, self](asio::error_code ec, const tcp::endpoint&) {
                            if (!ec) {
                                mConnected = true;
                                Utils::log("Client connected.");
                                mHandler = std::make_shared<MessageHandler>(std::move(mSocket));
                                mHandler->onMessage([this](AnyMessageType type, const std::string& data) {
                                    mHandleMessage(static_cast<ServerMessageType>(MessageHandler::toUint8(type)), data);
                                });
                                // Close only after peer FIN to avoid RSTs
                                mHandler->onDisconnect([this](const asio::error_code& ec) {
                                    asio::error_code ec2;
                                    if (mHandler) {
                                        mHandler->socket().close(ec2);
                                    }
                                    mConnected = false;
                                    Utils::log(std::string("Server disconnected: ") + ec.message());
                                });
                                mHandler->start();
                                
                                // Init connection handshake
                                Utils::log("Sending handshake init to server.");

                                // Check if hostname or IPv4/IPv6
                                try {
                                    asio::ip::make_address(mHost);
                                    self->mIsHostDomain = false; // IPv4 or IPv6 literal
                                } catch (const asio::system_error&) {
                                    self->mIsHostDomain = true;  // hostname / domain
                                }


                                std::vector<uint8_t> payload;
                                payload.reserve(1 + crypto_box_PUBLICKEYBYTES);
                                payload.push_back(Utils::protocolVersion());
                                /*payload.insert(payload.end(),
                                    mLibSodiumWrapper->getXPublicKey(),
                                    mLibSodiumWrapper->getXPublicKey() + crypto_box_PUBLICKEYBYTES
                                );*/
                                payload.insert(payload.end(),
                                    mLibSodiumWrapper->getPublicKey(),
                                    mLibSodiumWrapper->getPublicKey() + crypto_sign_PUBLICKEYBYTES
                                );

                                mHandler->sendMessage(ClientMessageType::HANDSHAKE_INIT, Utils::uint8ArrayToString(payload.data(), payload.size()));
                            
                                mStartHeartbeat();
                            } else {
                                if (!NetHelper::isExpectedDisconnect(ec)) {
                                    Utils::error("Client connect failed: " + ec.message());
                                }
                            }
                        });
                } else {
                    Utils::error("Client resolve failed: " + ec.message());
                }
            });
    }

    void TCPClient::sendMessage(ClientMessageType type, const std::string& data) {
        if (!mConnected) {
            Utils::error("Cannot send message, client not connected.");
            return;
        }

        if (mHandler) {
            asio::post(mHandler->socket().get_executor(), [self = shared_from_this(), type, data]() {
                self->mHandler->sendMessage(type, data);
            });
        }
    }

    void TCPClient::disconnect(bool echo) {
        if (mConnected && mHandler) {
            if (echo) {
                mHandler->sendMessage(ClientMessageType::GRACEFUL_DISCONNECT, "Goodbye");
            }

            asio::error_code ec;
            mHeartbeatTimer.cancel();

            // Half-close: stop sending, keep reading until peer FIN
            mHandler->socket().shutdown(tcp::socket::shutdown_send, ec);
            if (ec) {
                Utils::error("Error during socket shutdown: " + ec.message());
            }

            // Do not close immediately; rely on onDisconnect to finalize
            Utils::log("Client initiated graceful disconnect (half-close).");
        }
    }

    bool TCPClient::isHandshakeComplete() const {
        return mHandshakeComplete;
    }

    bool TCPClient::isConnected() const {
        return mConnected;
    }

    void TCPClient::mStartHeartbeat() {
        auto self = shared_from_this();
        mHeartbeatTimer.expires_after(std::chrono::seconds(5));
        mHeartbeatTimer.async_wait([this, self](const asio::error_code& ec) {
            if (ec == asio::error::operation_aborted) {
                return; // Timer was cancelled
            }

            auto now = std::chrono::steady_clock::now();
            auto elapsed = std::chrono::duration_cast<std::chrono::seconds>(now - self->mLastHeartbeatReceived).count();

            if (elapsed >= 15) { // 3 missed heartbeats
                Utils::error("Missed 3 heartbeats. I think the other party might have died! Disconnecting.");
                
                // Close sockets forcefully, server is dead
                asio::error_code ec;
                mHandler->socket().shutdown(tcp::socket::shutdown_both, ec);
                mHandler->socket().close(ec);
                mConnected = false;

                mGlobalKeyRef = nullptr;
                if (mSessionIDRef) {
                    *mSessionIDRef = 0;
                }

                return;
            }

            self->sendMessage(ClientMessageType::HEARTBEAT);
            Utils::debug("Sent HEARTBEAT to server.");
            self->mLastHeartbeatSent = std::chrono::steady_clock::now();
            
            self->mStartHeartbeat(); // Recursive
        });
    }

    void TCPClient::mHandleMessage(ServerMessageType type, const std::string& data) {
        switch (type) {
            case ServerMessageType::HANDSHAKE_IDENTIFY: {
                    Utils::log("Received server identity: " + data);
                    std::memcpy(mServerPublicKey, data.data(), std::min(data.size(), sizeof(mServerPublicKey)));

                    // Verify pubkey against whitelisted_keys
                    std::vector<std::string> whitelistedKeys = Utils::getWhitelistedKeys();
                    if (std::find(whitelistedKeys.begin(), whitelistedKeys.end(), Utils::bytesToHexString(mServerPublicKey, 32)) == whitelistedKeys.end()) { // Key verification is handled in later steps of the handshake
                        if (!mInsecureMode) {
                            Utils::error("Server public key not in whitelisted_keys. Terminating connection.");
                            disconnect();
                            return;
                        }

                        Utils::warn("Server public key not in whitelisted_keys, but continuing due to insecure mode.");
                    }

                    // Generate and send challenge
                    Utils::log("Sending challenge to server.");
                    mSubmittedChallenge = Utils::LibSodiumWrapper::generateRandom256Bit(); // Temporarily store the challenge to verify later
                    mHandler->sendMessage(ClientMessageType::HANDSHAKE_CHALLENGE, Utils::uint8ArrayToString(mSubmittedChallenge));
                }

                break;
            case ServerMessageType::HANDSHAKE_CHALLENGE_RESPONSE:
                Utils::log("Received challenge response from server.");
                {
                    // Verify the signature
                    Signature sig{};
                    std::memcpy(sig.data(), data.data(), std::min(data.size(), sig.size()));
                    if (Utils::LibSodiumWrapper::verifyMessage(mSubmittedChallenge.data(), mSubmittedChallenge.size(), sig, mServerPublicKey)) {
                        Utils::log("Challenge response verified successfully.");
                        
                        // Convert the server's public key to Curve25519 for encryption
                        AsymPublicKey serverXPubKey{};
                        int r = crypto_sign_ed25519_pk_to_curve25519(serverXPubKey.data(), mServerPublicKey);
                        if (r != 0) {
                            Utils::error("Failed to convert server signing key to encryption key! Killing connection.");
                            disconnect();
                            return;
                        }

                        // Generate AES key and send confirmation
                        mConnectionAESKey = Utils::LibSodiumWrapper::generateRandom256Bit();
                        if (mGlobalKeyRef) { // Copy to the global reference
                            std::copy(mConnectionAESKey.begin(), mConnectionAESKey.end(), mGlobalKeyRef->begin());
                        }
                        AsymNonce nonce{};
                        randombytes_buf(nonce.data(), nonce.size());

                        // TODO: This is pretty redundant, it should return the required type directly
                        std::array<uint8_t, 32> arrayPrivateKey;
                        std::copy(mLibSodiumWrapper->getXPrivateKey(),
                                  mLibSodiumWrapper->getXPrivateKey() + 32,
                                  arrayPrivateKey.begin());

                        std::vector<uint8_t> encr = Utils::LibSodiumWrapper::encryptAsymmetric(
                            mConnectionAESKey.data(), mConnectionAESKey.size(),
                            nonce,
                            serverXPubKey,
                            arrayPrivateKey
                        );

                        std::vector<uint8_t> payload;
                        payload.reserve(nonce.size() + encr.size());
                        payload.insert(payload.end(), nonce.begin(), nonce.end());
                        payload.insert(payload.end(), encr.begin(), encr.end());

                        mHandler->sendMessage(ClientMessageType::HANDSHAKE_EXCHANGE_KEY, Utils::uint8ArrayToString(payload.data(), payload.size()));
                    } else {
                        Utils::error("Challenge response verification failed. Terminating connection.");
                        disconnect();
                    }
                }

                break;
            case ServerMessageType::HANDSHAKE_EXCHANGE_KEY_CONFIRM:
                Utils::log("Received handshake exchange key confirmation from server.");
                // Decrypt the session ID using the established AES key
                {
                    Nonce symNonce{}; // All zeros
                    std::vector<uint8_t> ciphertext(data.begin(), data.end());
                    std::vector<uint8_t> decrypted = Utils::LibSodiumWrapper::decryptMessage(
                        ciphertext.data(), ciphertext.size(),
                        mConnectionAESKey, symNonce
                    );

                    if (decrypted.size() != sizeof(mConnectionSessionID) + sizeof(Protocol::TunConfig)) {
                        Utils::error("Decrypted config has invalid size. Terminating connection.");
                        disconnect();
                        return;
                    }

                    std::memcpy(&mConnectionSessionID, decrypted.data(), sizeof(mConnectionSessionID));
                    std::memcpy(&mTunConfig, decrypted.data() + sizeof(mConnectionSessionID), sizeof(Protocol::TunConfig));

                    mConnectionSessionID = Utils::cbe64toh(mConnectionSessionID);

                    Utils::log("Connection established with Session ID: " + std::to_string(mConnectionSessionID));
                
                    if (mSessionIDRef) { // Copy to the global reference
                        *mSessionIDRef = mConnectionSessionID;
                    }

                    uint32_t clientIP = ntohl(mTunConfig.clientIP);
                    uint32_t serverIP = ntohl(mTunConfig.serverIP);
                    uint8_t prefixLen = mTunConfig.prefixLength;
                    uint16_t mtu = mTunConfig.mtu;

                    if (mTun) {
                        mTun->configureIP(clientIP, serverIP, prefixLen, mtu);
                    }

                    mHandshakeComplete = true;
                }
            
                break;
            case ServerMessageType::HEARTBEAT:
                Utils::debug("Received HEARTBEAT from server.");
                mHandler->sendMessage(ClientMessageType::HEARTBEAT_ACK, ""); // Send ACK
                break;
            case ServerMessageType::HEARTBEAT_ACK:
                Utils::debug("Received HEARTBEAT_ACK from server.");
                mLastHeartbeatReceived = std::chrono::steady_clock::now();
                mMissedHeartbeats = 0; // Reset missed heartbeat count
                break;
            case ServerMessageType::GRACEFUL_DISCONNECT:
                Utils::log("Server is disconnecting: " + data);
                if (mConnected) { // Prevent Recursion
                    disconnect(false);
                }
                break;
            case ServerMessageType::KILL_CONNECTION:
                Utils::warn("Server is killing the connection: " + data);
                if (mConnected) {
                    disconnect(false);
                }
                break;
            default:
                Utils::log("Received unknown message type from server.");
                break;
        }
    }
}