#define ASIO_STANDALONE

#include <base64.hpp>
#include <chrono>
#include <fstream>
#include <functional>
#include <json.hpp>
#include <mutex>
#include <random>
#include <thread>
#include <websocketpp/config/asio_no_tls.hpp>
#include <websocketpp/server.hpp>

#include "cardset.h"
#include "database.h"
#include "imagestore.h"
#include "wizard.h"

namespace {

using socket_type = websocketpp::server<websocketpp::config::asio>;

class server {
 public:
  server(const cardset& cards, const imagestore& images, int port,
         std::mt19937& rng)
      : cards_(cards), images_(images), port_(port), rng_(rng) {
    socket_.init_asio();

    socket_.set_message_handler([this](websocketpp::connection_hdl connection,
                                       socket_type::message_ptr message) {
      on_message(connection, message);
    });
  }

  void run() {
    socket_.listen(port_);
    socket_.start_accept();

    asio::post(std::bind(&server::cleanup_thread, this));

    std::cout << "Listening on port " << port_ << "..." << std::endl;

    socket_.run();
  }

 private:
  void on_message(websocketpp::connection_hdl connection,
                  socket_type::message_ptr message) {
    nlohmann::json msgJson;

    try {
      msgJson = nlohmann::json::parse(message->get_payload());
    } catch (const std::exception& ex) {
      std::string response = R"(
        {
          "type": "error",
          "msg": "Could not parse message."
        }
      )";

      socket_.send(connection, response,
                   websocketpp::frame::opcode::value::TEXT);
      return;
    }

    std::string cmd = msgJson["cmd"];
    if (cmd == "generate") {
      cmd_generate(connection, msgJson["text"]);
    } else if (cmd == "check") {
      cmd_check(connection, msgJson["token"]);
    } else {
      std::string response = R"(
        {
          "type": "error",
          "msg": "No command in message."
        }
      )";

      socket_.send(connection, response,
                   websocketpp::frame::opcode::value::TEXT);
      return;
    }
  }

  void cmd_generate(websocketpp::connection_hdl connection, std::string text) {
    std::string token;

    {
      std::lock_guard rng_guard(rng_mutex_);
      token = database_.create(rng_);
    }

    nlohmann::json tokenMsg;
    tokenMsg["type"] = "token";
    tokenMsg["token"] = token;
    socket_.send(connection, tokenMsg.dump(),
                 websocketpp::frame::opcode::value::TEXT);

    database_.subscribe(token, [this, connection](const std::string& msg) {
      socket_.send(connection, msg, websocketpp::frame::opcode::value::TEXT);
    });

    asio::post(std::bind(&server::generate_thread, this, token, text));
  }

  void cmd_check(websocketpp::connection_hdl connection, std::string token) {
    bool failed = false;

    try {
      database_.subscribe(token, [this, connection](const std::string& msg) {
        socket_.send(connection, msg, websocketpp::frame::opcode::value::TEXT);
      });

      if (!database_.is_done(token)) {
        return;
      }

      std::string result = database_.getResult(token);
      nlohmann::json resultMsg;
      if (result.empty()) {
        resultMsg["type"] = "error";
        resultMsg["msg"] = "Unknown error occurred.";
      } else {
        resultMsg["type"] = "result";
        resultMsg["image"] = database_.getResult(token);
        resultMsg["msg"] = "Success!";
      }

      socket_.send(connection, resultMsg.dump(),
                   websocketpp::frame::opcode::value::TEXT);
    } catch (const std::exception& ex) {
      failed = true;
    }

    if (failed) {
      try {
        socket_.send(connection, R"(
          {
            "type": "error",
            "msg": "Error retrieving request."
          })",
                     websocketpp::frame::opcode::value::TEXT);
      } catch (const std::exception& ex) {
        // Well, okay
        std::cout << ex.what() << std::endl;
      }
    }
  }

  void generate_thread(std::string token, std::string text) {
    std::unique_ptr<wizard> generator;

    {
      std::lock_guard rng_guard(rng_mutex_);
      generator = std::make_unique<wizard>(cards_, images_, text, rng_);
    }

    try {
      generator->set_status_callback([this, token](const std::string& status) {
        nlohmann::json msg;
        msg["type"] = "status";
        msg["msg"] = status;

        database_.post(token, msg.dump());
      });

      Magick::Image resultImage = generator->run();
      Magick::Blob resultBlob;
      resultImage.write(&resultBlob);

      std::string resultBytes((const char*)resultBlob.data(),
                              resultBlob.length());
      std::string resultEncoded = base64::to_base64(resultBytes);

      database_.setResult(token, resultEncoded);

      nlohmann::json resultMsg;
      resultMsg["type"] = "result";
      resultMsg["image"] = resultEncoded;
      resultMsg["msg"] = "Success!";

      database_.post(token, resultMsg.dump());
    } catch (const std::exception& ex) {
      nlohmann::json response;
      response["type"] = "error";
      response["msg"] =
          std::string("Error generating card (") + ex.what() + ")";

      database_.post(token, response.dump());
    }

    database_.mark_done(token);
  }

  void cleanup_thread() {
    for (;;) {
      // sleep
    }
  }

  const cardset& cards_;
  const imagestore& images_;
  const int port_;

  std::mutex rng_mutex_;
  std::mt19937& rng_;

  socket_type socket_;
  database database_;
};

}  // namespace

int main(int argc, char** argv) {
  Magick::InitializeMagick(nullptr);

  std::random_device randomDevice;
  std::mt19937 rng(5);  // randomDevice());

  if (argc != 2) {
    std::cout << "usage: wizard_server [configfile]" << std::endl;
    return -1;
  }

  std::ifstream config_file(argv[1]);
  nlohmann::json config_data = nlohmann::json::parse(config_file);

  cardset cards(config_data["cards_path"]);
  imagestore images(config_data["cache_path"]);
  int port = config_data["port"];

  server app(cards, images, port, rng);
  app.run();
}