#include <catch2/catch.hpp>
#include <crypto/crypto.hpp>
#include <crypto/crypto_libsodium.hpp>
#include <string_view>

#include <router_contact.hpp>
#include <iwp/iwp.hpp>
#include <util/meta/memfn.hpp>
#include <messages/link_message_parser.hpp>
#include <messages/discard.hpp>
#include <util/time.hpp>

#include <net/net_if.hpp>

#undef LOG_TAG
#define LOG_TAG __FILE__

namespace iwp = llarp::iwp;
namespace util = llarp::util;

/// make an iwp link
template <bool inbound, typename... Args_t>
static llarp::LinkLayer_ptr
make_link(Args_t... args)
{
  if (inbound)
    return iwp::NewInboundLink(args...);
  else
    return iwp::NewOutboundLink(args...);
}
using Logic_ptr = std::shared_ptr<llarp::Logic>;

/// a single iwp link with associated keys and members to make unit tests work
struct IWPLinkContext
{
  llarp::RouterContact rc;
  llarp::IpAddress localAddr;
  llarp::LinkLayer_ptr link;
  std::shared_ptr<llarp::KeyManager> keyManager;
  llarp::LinkMessageParser m_Parser;
  llarp_ev_loop_ptr m_Loop;
  /// is the test done on this context ?
  bool gucci = false;

  IWPLinkContext(std::string_view addr, llarp_ev_loop_ptr loop)
      : localAddr{std::move(addr)}
      , keyManager{std::make_shared<llarp::KeyManager>()}
      , m_Parser{nullptr}
      , m_Loop{std::move(loop)}
  {
    // generate keys
    llarp::CryptoManager::instance()->identity_keygen(keyManager->identityKey);
    llarp::CryptoManager::instance()->encryption_keygen(keyManager->encryptionKey);
    llarp::CryptoManager::instance()->encryption_keygen(keyManager->transportKey);

    // set keys in rc
    rc.pubkey = keyManager->identityKey.toPublic();
    rc.enckey = keyManager->encryptionKey.toPublic();
  }

  template <typename Func_t>
  void
  Call(Func_t work)
  {
    m_Loop->call_soon(std::move(work));
  }

  bool
  HandleMessage(llarp::ILinkSession* from, const llarp_buffer_t& buf)
  {
    return m_Parser.ProcessFrom(from, buf);
  }

  /// initialize link
  template <bool inbound>
  void
  InitLink(std::function<void(llarp::ILinkSession*)> established)
  {
    link = make_link<inbound>(
        keyManager,
        m_Loop,
        // getrc
        [&]() -> const llarp::RouterContact& { return rc; },
        // link message handler
        util::memFn(&IWPLinkContext::HandleMessage, this),
        // sign buffer
        [&](llarp::Signature& sig, const llarp_buffer_t& buf) {
          REQUIRE(llarp::CryptoManager::instance()->sign(sig, keyManager->identityKey, buf));
          return true;
        },
        // before connect
        nullptr,
        // established handler
        [established](llarp::ILinkSession* s, bool linkIsInbound) {
          REQUIRE(s != nullptr);
          REQUIRE(inbound == linkIsInbound);
          established(s);
          return true;
        },
        // renegotiate handler
        [](llarp::RouterContact newrc, llarp::RouterContact oldrc) {
          REQUIRE(newrc.pubkey == oldrc.pubkey);
          return true;
        },
        // timeout handler
        [&](llarp::ILinkSession*) {
          llarp_ev_loop_stop(m_Loop);
          FAIL("session timeout");
        },
        // session closed handler
        [](llarp::RouterID) {},
        // pump done handler
        []() {},
        // do work function
        [l = m_Loop](llarp::Work_t work) { l->call_soon(work); });
    REQUIRE(link->Configure(
        m_Loop, llarp::net::LoopbackInterfaceName(), AF_INET, *localAddr.getPort()));

    if (inbound)
    {
      // only add address info on the recipiant's rc
      rc.addrs.emplace_back();
      REQUIRE(link->GetOurAddressInfo(rc.addrs.back()));
    }
    // sign rc
    REQUIRE(rc.Sign(keyManager->identityKey));
    REQUIRE(keyManager != nullptr);
  }
};

using Context_ptr = std::shared_ptr<IWPLinkContext>;

/// run an iwp unit test after setup
/// call take 2 parameters, test and a timeout
///
/// test is a callable that takes 5 arguments:
/// 0) std::function<Logic_ptr(void)> that starts the iwp links and gives a logic to call with
/// 1) std::function<void(void)> that ends the unit test if we are done
/// 2) std::function<void(void)> that ends the unit test right now as a success
/// 3) client iwp link context (shared_ptr)
/// 4) relay iwp link context (shared_ptr)
///
/// timeout is a std::chrono::duration that tells the driver how long to run the unit test for
/// before it should assume failure of unit test
template <typename Func_t, typename Duration_t = std::chrono::milliseconds>
void
RunIWPTest(Func_t test, Duration_t timeout = 10s)
{
  // shut up logs
  llarp::LogSilencer shutup;
  // set up event loop
  auto logic = std::make_shared<llarp::Logic>();
  auto loop = llarp_make_ev_loop();
  loop->set_logic(logic);

  llarp::LogContext::Instance().Initialize(
      llarp::eLogDebug, llarp::LogType::File, "stdout", "unit test", [loop](auto work) {
        loop->call_soon(work);
      });

  // turn off bogon blocking
  auto oldBlockBogons = llarp::RouterContact::BlockBogons;
  llarp::RouterContact::BlockBogons = false;

  // set up cryptography
  llarp::sodium::CryptoLibSodium crypto{};
  llarp::CryptoManager manager{&crypto};

  // set up client
  auto initiator = std::make_shared<IWPLinkContext>("127.0.0.1:3001", loop);
  // set up server
  auto recipiant = std::make_shared<IWPLinkContext>("127.0.0.1:3002", loop);

  // function for ending unit test on success
  auto endIfDone = [initiator, recipiant, loop, logic]() {
    if (initiator->gucci and recipiant->gucci)
    {
      LogicCall(logic, [loop]() { llarp_ev_loop_stop(loop); });
    }
  };
  // function to start test and give logic to unit test
  auto start = [initiator, recipiant, logic]() {
    REQUIRE(initiator->link->Start(logic));
    REQUIRE(recipiant->link->Start(logic));
    return logic;
  };

  // function to end test immediately
  auto endTest = [logic, loop]() { LogicCall(logic, [loop]() { llarp_ev_loop_stop(loop); }); };

  loop->call_after_delay(
      std::chrono::duration_cast<llarp_time_t>(timeout), []() { FAIL("test timeout"); });
  test(start, endIfDone, endTest, initiator, recipiant);
  llarp_ev_loop_run_single_process(loop, logic);
  llarp::RouterContact::BlockBogons = oldBlockBogons;
}

/// ensure clients can connect to relays
TEST_CASE("IWP handshake", "[iwp]")
{
  RunIWPTest([](std::function<Logic_ptr(void)> start,
                std::function<void(void)> endIfDone,
                [[maybe_unused]] std::function<void(void)> endTestNow,
                Context_ptr alice,
                Context_ptr bob) {
    // set up initiator
    alice->InitLink<false>([=](auto remote) {
      REQUIRE(remote->GetRemoteRC() == bob->rc);
      alice->gucci = true;
      endIfDone();
    });
    // set up recipiant
    bob->InitLink<true>([=](auto remote) {
      REQUIRE(remote->GetRemoteRC() == alice->rc);
      bob->gucci = true;
      endIfDone();
    });
    // start unit test
    auto logic = start();
    // try establishing a session
    LogicCall(logic, [link = alice->link, rc = bob->rc]() { REQUIRE(link->TryEstablishTo(rc)); });
  });
}

/// ensure relays cannot connect to clients
TEST_CASE("IWP handshake reverse", "[iwp]")
{
  RunIWPTest([](std::function<Logic_ptr(void)> start,
                [[maybe_unused]] std::function<void(void)> endIfDone,
                std::function<void(void)> endTestNow,
                Context_ptr alice,
                Context_ptr bob) {
    alice->InitLink<false>([](auto) {});
    bob->InitLink<true>([](auto) {});
    // start unit test
    auto logic = start();
    // try establishing a session in the wrong direction
    LogicCall(logic, [logic, link = bob->link, rc = alice->rc, endTestNow]() {
      REQUIRE(not link->TryEstablishTo(rc));
      endTestNow();
    });
  });
}

/// ensure iwp can send messages between sessions
TEST_CASE("IWP send messages", "[iwp]")
{
  int aliceNumSent = 0;
  int bobNumSent = 0;
  RunIWPTest([&aliceNumSent, &bobNumSent](std::function<Logic_ptr(void)> start,
                std::function<void(void)> endIfDone,
                std::function<void(void)> endTestNow,
                Context_ptr alice,
                Context_ptr bob) {
    constexpr int numSend = 128;
    // when alice makes a session to bob send `aliceNumSend` messages to him
    alice->InitLink<false>([endIfDone, alice, &aliceNumSent](auto session) {
      for (auto index = 0; index < numSend; index++)
      {
        alice->Call([session, endIfDone, alice, &aliceNumSent]() {
          // generate a discard message that is 512 bytes long
          llarp::DiscardMessage msg;
          std::vector<byte_t> msgBuff(512);
          llarp_buffer_t buf(msgBuff);
          // add random padding
          llarp::CryptoManager::instance()->randomize(buf);
          // encode the discard message
          msg.BEncode(&buf);
          // send the message
          session->SendMessageBuffer(msgBuff, [endIfDone, alice, &aliceNumSent](auto status) {
            if (status == llarp::ILinkSession::DeliveryStatus::eDeliverySuccess)
            {
              // on successful transmit increment the number we sent
              aliceNumSent++;
            }
            // if we sent all the messages sucessfully we end the unit test
            alice->gucci = aliceNumSent == numSend;
            endIfDone();
          });
        });
      }
    });
    bob->InitLink<true>([endIfDone, bob, &bobNumSent](auto session) {
      for (auto index = 0; index < numSend; index++)
      {
        bob->Call([session, endIfDone, bob, &bobNumSent]() {
          // generate a discard message that is 512 bytes long
          llarp::DiscardMessage msg;
          std::vector<byte_t> msgBuff(512);
          llarp_buffer_t buf(msgBuff);
          // add random padding
          llarp::CryptoManager::instance()->randomize(buf);
          // encode the discard message
          msg.BEncode(&buf);
          // send the message
          session->SendMessageBuffer(msgBuff, [endIfDone, bob, &bobNumSent](auto status) {
            if (status == llarp::ILinkSession::DeliveryStatus::eDeliverySuccess)
            {
              // on successful transmit increment the number we sent
              bobNumSent++;
            }
            // if we sent all the messages sucessfully we end the unit test
            bob->gucci = bobNumSent == numSend;
            endIfDone();
          });
        });
      }
    });
    // start unit test
    auto logic = start();
    // try establishing a session from alice to bob
    LogicCall(logic, [logic, link = alice->link, rc = bob->rc, endTestNow]() {
      REQUIRE(link->TryEstablishTo(rc));
    });
  });
}