/*
 * Copyright (C) 2025 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <chrono>
#include <iostream>
#include <cstring>
#include <vector>
#include <thread>
#include <chrono>
#include <cstring>
#include <unistd.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <poll.h>
#include <unistd.h>
#include <random>
#include <future>
#include <thread>

#include <android-base/file.h>
#include <android-base/logging.h>
#include <android-base/scopeguard.h>
#include <IOUringSocketHandler/IOUringSocketHandler.h>

#include <gtest/gtest.h>

// Test all combinations of queue_depth and messages
struct TestParam {
    int queue_depth;
    int numMessages;
    int messageSize;
};

class IOUringSocketHandlerTest : public ::testing::TestWithParam<TestParam> {
public:
    bool IsIouringSupported() {
        return IOUringSocketHandler::IsIouringSupported();
    }
    void SendMsg(int sock_send, const bool non_block);

protected:
    void SetUp() override {
    }
    void TearDown() override {
        close(sock_recv_);
        close(sock_send_);
    }
    void ReceiveThreaduring(int sock_recv);
    bool CreateServerSocket();
    std::unique_ptr<IOUringSocketHandler> handler_;
    void InitializeHandler(int socket_fd = 1);
    // Default queue depth
    int queue_depth_ = 1;
    int sock_recv_;
    int sock_send_;
    int socks[2];
    std::vector<std::string> sent_messages; // Store sent messages for comparison
};

bool IOUringSocketHandlerTest::CreateServerSocket() {
    if (socketpair(AF_UNIX, SOCK_DGRAM, 0, socks) < 0) {
        PLOG(ERROR) << "socketpair failed";
        return false;
    }

    int on = 1;
    if (setsockopt(socks[0], SOL_SOCKET, SO_PASSCRED, &on, sizeof(on)) != 0) {
        PLOG(ERROR) << "setsockopt SO_PASSCRED failed for socks[0]";
        close(socks[0]);
        close(socks[1]);
        return false;
    }

    if (setsockopt(socks[1], SOL_SOCKET, SO_PASSCRED, &on, sizeof(on)) != 0) {
        PLOG(ERROR) << "setsockopt SO_PASSCRED failed for socks[1]";
        close(socks[0]);
        close(socks[1]);
        return false;
    }

    sock_recv_ = socks[0];
    sock_send_ = socks[1];
    return true;
}

// Function to generate a random string
static std::string generateRandomString(size_t length) {
    static const char charset[] =
        "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> dis(0, sizeof(charset) - 2);

    std::string str(length, 0);
    for (size_t i = 0; i < length; ++i) {
        str[i] = charset[dis(gen)];
    }
    return str;
}

void IOUringSocketHandlerTest::InitializeHandler(int socket_fd) {
    handler_ = std::make_unique<IOUringSocketHandler>(socket_fd);
}

TEST_F(IOUringSocketHandlerTest, SetupIoUring) {
    if (!IsIouringSupported()) {
        GTEST_SKIP() << "io_uring not supported. Skipping Test.";
    }
    InitializeHandler();
    EXPECT_TRUE(handler_->SetupIoUring(queue_depth_));
}

TEST_F(IOUringSocketHandlerTest, AllocateAndRegisterBuffers) {
    if (!IsIouringSupported()) {
        GTEST_SKIP() << "io_uring not supported. Skipping Test.";
    }
    InitializeHandler();
    EXPECT_TRUE(handler_->SetupIoUring(queue_depth_));
    EXPECT_TRUE(handler_->AllocateAndRegisterBuffers(8, 4096));
}

TEST_F(IOUringSocketHandlerTest, MultipleAllocateAndRegisterBuffers) {
    if (!IsIouringSupported()) {
        GTEST_SKIP() << "io_uring not supported. Skipping Test.";
    }
    InitializeHandler();

    EXPECT_TRUE(handler_->SetupIoUring(queue_depth_));

    EXPECT_TRUE(handler_->AllocateAndRegisterBuffers(4, 4096));
    handler_->DeRegisterBuffers();

    EXPECT_TRUE(handler_->AllocateAndRegisterBuffers(2, 1024*1024L));
    handler_->DeRegisterBuffers();

    EXPECT_TRUE(handler_->AllocateAndRegisterBuffers(32, 1024));
    handler_->DeRegisterBuffers();

    // num_buffers should be power of 2
    EXPECT_FALSE(handler_->AllocateAndRegisterBuffers(5, 4096));
}

TEST_F(IOUringSocketHandlerTest, TestStopAPI) {
    if (!IsIouringSupported()) {
        GTEST_SKIP() << "io_uring not supported. Skipping Test.";
    }
    ASSERT_TRUE(CreateServerSocket());
    InitializeHandler(sock_recv_);
    ASSERT_TRUE(handler_->SetupIoUring(16));
    handler_->AllocateAndRegisterBuffers(16, 4096);
    ASSERT_TRUE(handler_->EnqueueMultishotRecvmsg());

    std::atomic<bool> thread_started{false};
    std::atomic<bool> thread_exited{false};
    std::thread receive_thread([this, &thread_started, &thread_exited]() {
        thread_started = true;
        struct ucred* cred = nullptr;
        void* this_recv = nullptr;
        size_t len = 0;
        // This should block until Stop is called
        handler_->ReceiveData(&this_recv, len, &cred);

        // After Stop is called, these should be nullptr/0
        EXPECT_EQ(this_recv, nullptr);
        EXPECT_EQ(len, 0);
        EXPECT_EQ(cred, nullptr);

        thread_exited = true;
    });

    while (!thread_started) std::this_thread::yield();
    std::this_thread::sleep_for(std::chrono::milliseconds(50));

    handler_->Stop();
    receive_thread.join();
    EXPECT_TRUE(thread_exited);
}

void IOUringSocketHandlerTest::SendMsg(int sock_send, const bool non_block) {
   const TestParam params = GetParam();
   for (int i = 0; i < params.numMessages; ++i) {
     std::string message = generateRandomString(params.messageSize);

     sent_messages.push_back(message);
     struct ucred cred;
     memset(&cred, 0, sizeof(cred));
     cred.pid = getpid();
     cred.uid = getuid();
     cred.gid = getgid();

     struct iovec iov_send;
     iov_send.iov_base = const_cast<char*>(message.data());
     iov_send.iov_len = params.messageSize;

     struct msghdr msg_send;
     memset(&msg_send, 0, sizeof(msg_send));
     msg_send.msg_iov = &iov_send;
     msg_send.msg_iovlen = 1;

     char control_buffer_send[CMSG_SPACE(sizeof(cred))];
     memset(control_buffer_send, 0, sizeof(control_buffer_send));
     msg_send.msg_control = control_buffer_send;
     msg_send.msg_controllen = sizeof(control_buffer_send);

     struct cmsghdr* cmsg_send = CMSG_FIRSTHDR(&msg_send);
     cmsg_send->cmsg_level = SOL_SOCKET;
     cmsg_send->cmsg_type = SCM_CREDENTIALS;
     cmsg_send->cmsg_len = CMSG_LEN(sizeof(cred));
     memcpy(CMSG_DATA(cmsg_send), &cred, sizeof(cred));

     int flags = 0;
     if (non_block) {
        flags = MSG_DONTWAIT;
     }
     ssize_t sent_bytes;
        while (true) {
            sent_bytes = sendmsg(sock_send, &msg_send, flags);
            if (sent_bytes >= 0) {
                break; // Success
            }
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                // Try again
                continue;
            } else {
                perror("sendmsg failed");
                return;
            }
        }
    }
}

void IOUringSocketHandlerTest::ReceiveThreaduring(int sock_recv) {
    std::unique_ptr<IOUringSocketHandler> uring_listener;
    uring_listener = std::make_unique<IOUringSocketHandler>(sock_recv);
    const TestParam params = GetParam();
    ASSERT_TRUE(uring_listener->SetupIoUring(params.queue_depth));
    uring_listener->AllocateAndRegisterBuffers(
        params.queue_depth, params.messageSize);

    ASSERT_TRUE(uring_listener->EnqueueMultishotRecvmsg());

    long long received_messages = 0;
    int index = 0;
    while (received_messages < params.numMessages) {
        struct ucred* cred = nullptr;
        void* this_recv = nullptr;
        size_t len = 0;
        uring_listener->ReceiveData(&this_recv, len, &cred);
        // Release the buffer from here onwards
        {
            auto scope_guard =
                android::base::make_scope_guard([&uring_listener]() -> void {
                  uring_listener->ReleaseBuffer(); });

            if (len <= 0) {
                continue;
            }
            received_messages++;
            char* char_ptr = static_cast<char*>(this_recv);
            std::string payload_string(char_ptr, len);
            std::string orig_string = sent_messages[index];
            // Compare payload data
            EXPECT_EQ(payload_string, orig_string);
            // Verify credentials
            EXPECT_EQ(cred->uid, getuid());
            EXPECT_EQ(cred->gid, getgid());
            EXPECT_EQ(cred->pid, getpid());
            index += 1;
        }
    }
}

TEST_P(IOUringSocketHandlerTest, RecvmsgDataIntegrity) {
  if (!IsIouringSupported()) {
      GTEST_SKIP() << "io_uring not supported. Skipping Test.";
  }
  ASSERT_TRUE(CreateServerSocket());
  std::vector<std::thread> send_threads;
  send_threads.emplace_back([this](){ SendMsg(sock_send_, false); });
  ReceiveThreaduring(sock_recv_);
  for (auto& thread : send_threads) {
    thread.join();
  }
  close(sock_send_);
  close(sock_recv_);
}

TEST_P(IOUringSocketHandlerTest, RecvmsgDataIntegrityNonBlockingSend) {
  if (!IsIouringSupported()) {
      GTEST_SKIP() << "io_uring not supported. Skipping Test.";
  }
  ASSERT_TRUE(CreateServerSocket());
  int flags = fcntl(sock_send_, F_GETFL, 0);
  ASSERT_NE(fcntl(sock_send_, F_SETFL, flags | O_NONBLOCK), -1);

  std::vector<std::thread> send_threads;
  send_threads.emplace_back([this](){ SendMsg(sock_send_, true); });

  ReceiveThreaduring(sock_recv_);
  for (auto& thread : send_threads) {
    thread.join();
  }

  close(sock_send_);
  close(sock_recv_);
}

std::vector<TestParam> GetConfigs() {
  std::vector<TestParam> testParams;

  std::vector<int> queue_depth = {1, 8, 16, 32, 64, 128, 256, 512};
  std::vector<int> num_messages = {1, 100, 250, 500, 1000, 1500, 2000, 5000};
  std::vector<int> message_sizes = {1, 100, 520, 1024, 2042, 3168, 4068, 4096};

  // This will test 512 combinations
  for (auto message_size: message_sizes) {
    for (auto q_depth : queue_depth) {
      for (auto n_messages : num_messages) {
        TestParam param;
        param.queue_depth = q_depth;
        param.numMessages = n_messages;
        param.messageSize = message_size;
        testParams.push_back(std::move(param));
      }
    }
  }

  return testParams;
}

INSTANTIATE_TEST_SUITE_P(Io, IOUringSocketHandlerTest,
                         ::testing::ValuesIn(GetConfigs()));

int main(int argc, char** argv) {
    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}