#include <gtest/gtest.h>

#include <larra/client/client.hpp>
#include <larra/impl/mock_socket.hpp>
#include <larra/proxy.hpp>

using namespace larra::xmpp;
using boost::asio::ip::tcp;
namespace asio = boost::asio;

class ProxyTest : public ::testing::Test {
 protected:
  boost::asio::io_context io_context;
  larra::xmpp::impl::MockSocket mock_socket{io_context.get_executor()};
};

// Test 1: Connect via HTTP proxy with successful server response
TEST_F(ProxyTest, ConnectViaHttpProxy_SuccessfulResponse) {
  HttpProxy proxy{"proxy_host", 8080};

  std::string targetHost = "target_host";
  uint16_t targetPort = 80;

  std::string expectedRequest =
      std::format("CONNECT {}:{} HTTP/1.1\r\nHost: {}:{}\r\n\r\n", targetHost, targetPort, targetHost, targetPort);

  std::string proxyResponse = "HTTP/1.1 200 Connection established\r\n\r\n";

  mock_socket.AddReceivedData(proxyResponse);

  bool connectSuccessful = false;

  asio::co_spawn(
      io_context,
      [&]() -> asio::awaitable<void> {
        try {
          co_await client::impl::ConnectViaProxy(mock_socket, proxy, targetHost, targetPort);
          connectSuccessful = true;
        } catch(...) {
          connectSuccessful = false;
        }
      },
      asio::detached);

  io_context.run();

  std::string sentData = mock_socket.GetSentData();

  EXPECT_EQ(sentData, expectedRequest);
  EXPECT_TRUE(connectSuccessful);
}

// Test 2: Connect via SOCKS proxy
TEST(Socks5ProxyTest, ConnectViaProxy) {
  constexpr std::uint16_t kSocksPort = 1080;
  constexpr std::uint16_t kAvailableUdpBufferSpaceForSocks = 262;

  boost::asio::io_context io;
  auto executor = io.get_executor();

  larra::xmpp::impl::MockSocket socket{executor};

  std::string expectedServerResponse;
  expectedServerResponse += "\x05\x00";          // VER, METHOD
  expectedServerResponse += "\x05\x00\x00\x01";  // VER, REP, RSV, ATYP (IPv4)
  expectedServerResponse += "\x7F\x00\x00\x01";  // BND.ADDR (127.0.0.1)
  expectedServerResponse += "\x1F\x90";          // BND.PORT (8080)

  socket.AddReceivedData(expectedServerResponse);

  Socks5Proxy proxy{.hostname = "proxy.example.com", .port = kSocksPort};
  std::string targetHostname = "target.example.com";
  std::uint16_t targetPort = 80;

  boost::asio::co_spawn(
      executor,
      [&]() -> boost::asio::awaitable<void> {
        co_await client::impl::ConnectViaProxy(socket, proxy, targetHostname, targetPort);

        auto sentData = socket.GetSentData();

        std::string expectedGreeting = "\x05\x01\x00";

        std::array<std::uint8_t, kAvailableUdpBufferSpaceForSocks> expectedRequest{};
        std::size_t req_len = 0;

        expectedRequest[req_len++] = 0x05;  // VER
        expectedRequest[req_len++] = 0x01;  // CMD: CONNECT
        expectedRequest[req_len++] = 0x00;  // RSV
        expectedRequest[req_len++] = 0x03;  // ATYP: DOMAINNAME

        expectedRequest[req_len++] = static_cast<std::uint8_t>(targetHostname.size());  // domain length

        std::memcpy(&expectedRequest[req_len], targetHostname.data(), targetHostname.size());
        req_len += targetHostname.size();

        std::uint16_t networkOrderPort = htons(targetPort);
        expectedRequest[req_len++] = static_cast<std::uint8_t>((networkOrderPort >> 8) & 0xFF);
        expectedRequest[req_len++] = static_cast<std::uint8_t>(networkOrderPort & 0xFF);

        std::string expectedData = expectedGreeting;
        auto transformed_view = expectedRequest | std::views::take(req_len) | std::views::transform([](std::uint8_t byte) {
                                  return static_cast<char>(byte);
                                });

        expectedData.append(std::ranges::to<std::string>(transformed_view));
        EXPECT_EQ(sentData, expectedData);

        co_return;
      },
      boost::asio::detached);

  io.run();
}