#include <curl/curl.h>
#include <pistache/client.h>
#include <rapidjson/document.h>
#include <jwt-cpp/jwt.h>

#include <catch2/catch.hpp>
#include <daggy/Serialization.hpp>
#include <daggy/executors/task/NoopTaskExecutor.hpp>
#include <daggy/loggers/dag_run/OStreamLogger.hpp>
#include <daggyd/Server.hpp>
#include <filesystem>
#include <iostream>
#include <thread>

namespace rj = rapidjson;

using namespace daggy;

TEST_CASE("jwt_authentication", "[server_jwt]")
{
  std::stringstream ss;
  daggy::executors::task::NoopTaskExecutor executor;
  daggy::loggers::dag_run::OStreamLogger logger(ss);
  Pistache::Address listenSpec("localhost", Pistache::Port(0));

  const size_t nDAGRunners = 10, nWebThreads = 10;
  const std::string jwtSecret = "test-secret-key";

  daggy::daggyd::Server server(listenSpec, logger, executor, nDAGRunners,
                               "/dev/null");
  server.setJWTSecret(jwtSecret);
  server.init(nWebThreads);
  server.start();

  const std::string host    = "localhost:";
  const std::string baseURL = host + std::to_string(server.getPort());

  SECTION("Requests without JWT token should fail")
  {
    auto response = HTTP_REQUEST(baseURL + "/v1/dagruns");
    REQUIRE(response.code == HTTPCode::Unauthorized);
  }

  SECTION("Requests with invalid JWT token should fail")
  {
    std::map<std::string, std::string> headers = {
        {"Authorization", "Bearer invalid-token"}};
    auto response = HTTP_REQUEST_WITH_HEADERS(baseURL + "/v1/dagruns", "", "GET", headers);
    REQUIRE(response.code == HTTPCode::Unauthorized);
  }

  SECTION("Requests with valid JWT token should succeed")
  {
    auto token = jwt::create()
                     .set_issuer("daggy-test")
                     .set_type("JWT")
                     .sign(jwt::algorithm::hs256{jwtSecret});

    std::map<std::string, std::string> headers = {
        {"Authorization", "Bearer " + token}};
    auto response = HTTP_REQUEST_WITH_HEADERS(baseURL + "/v1/dagruns", "", "GET", headers);
    REQUIRE(response.code == HTTPCode::Ok);
  }

  SECTION("DAG submission with valid JWT token should work")
  {
    auto token = jwt::create()
                     .set_issuer("daggy-test")
                     .set_type("JWT")
                     .sign(jwt::algorithm::hs256{jwtSecret});

    std::string dagRun = R"({
            "tag": "jwt_test",
            "tasks": {
               "test_task": { "job": { "command": [ "/bin/echo", "hello" ], "environment": []} }
               }
        })";

    std::map<std::string, std::string> headers = {
        {"Authorization", "Bearer " + token}};
    auto response = HTTP_REQUEST_WITH_HEADERS(baseURL + "/v1/dagrun/", dagRun, "POST", headers);
    REQUIRE(response.code == HTTPCode::Ok);

    rj::Document doc;
    daggy::checkRJParse(doc.Parse(response.body.c_str()));
    REQUIRE(doc.IsObject());
    REQUIRE(doc.HasMember("runID"));
  }

  server.shutdown();
}

TEST_CASE("no_jwt_secret_allows_all", "[server_no_jwt]")
{
  std::stringstream ss;
  daggy::executors::task::NoopTaskExecutor executor;
  daggy::loggers::dag_run::OStreamLogger logger(ss);
  Pistache::Address listenSpec("localhost", Pistache::Port(0));

  const size_t nDAGRunners = 10, nWebThreads = 10;

  daggy::daggyd::Server server(listenSpec, logger, executor, nDAGRunners,
                               "/dev/null");
  server.init(nWebThreads);
  server.start();

  const std::string host    = "localhost:";
  const std::string baseURL = host + std::to_string(server.getPort());

  SECTION("Requests without JWT secret configured should work")
  {
    auto response = HTTP_REQUEST(baseURL + "/v1/dagruns");
    REQUIRE(response.code == HTTPCode::Ok);
  }

  server.shutdown();
}