#ifndef IPC_SESSION_H
#define IPC_SESSION_H

#include <mutex>
#include <memory> //std::enable_shared_from_this
#include <deque>

#if __GNUG__
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-variable"
#endif
#include "asio.hpp"
#if __GNUG__
#pragma GCC diagnostic pop
#endif

#include "ipc_session_holder.h"
#include "ipc_buffer.h"

namespace ipc {

class Session : public std::enable_shared_from_this<Session> {
  public:
    Session(asio::ip::tcp::socket socket, SessionHolder* session_holder, bool client) : socket_(std::move(socket)), session_holder_(session_holder), client_(client) { 
      // Add keepalive with timeout
      asio::socket_base::keep_alive option(true);
      socket_.set_option(option);

      // The timeout value of 5 seconds
      unsigned int timeout_milli = 5000;

      // This part is platform specific as ASIO doesn't allow setting timeout duration
      #if defined _WIN32 || defined WIN32 || defined OS_WIN64 || defined _WIN64 || defined WIN64 || defined WINNT
        // use windows-specific time
        int32_t timeout = timeout_milli;
        if (setsockopt(socket_.native(), SOL_SOCKET, SO_RCVTIMEO, (const char*)&timeout, sizeof(timeout)) < 0) {
          LOG_WARNING("Setting SO_RCVTIMEO failed! Timeout will not be set!");
        }
        if (setsockopt(socket_.native(), SOL_SOCKET, SO_SNDTIMEO, (const char*)&timeout, sizeof(timeout)) < 0) {
          LOG_WARNING("Setting SO_SNDTIMEO failed! Timeout will not be set!");
        }
        LOG_WARNING("On Windows keepalive timer might not work!");
      #else
        // POSIX
        // NOTE(harijs): If these options are set on acceptor() then the resulting socket would technically inherit these options, but that isn't cross-platform
        // NOTE(harijs): This doesn't seem to work
        struct timeval tv;
        tv.tv_sec  = timeout_milli / 1000;
        tv.tv_usec = timeout_milli % 1000;
        if (setsockopt(socket_.native(), SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)) < 0) {
          LOG_WARNING("Setting SO_RCVTIMEO failed! Timeout will not be set!");
        }
        if (setsockopt(socket_.native(), SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv)) < 0) {
          LOG_WARNING("Setting SO_SNDTIMEO failed! Timeout will not be set!");
        }
        // These are not portable - TCP_KEEPIDLE, TCP_KEEPINTVL and TCP_KEEPCNT, but could be set nonetheless
        // This is required for client to timeout if the client lost the connection
        int32_t timeout = timeout_milli;
        if (setsockopt(socket_.native(), SOL_TCP, TCP_USER_TIMEOUT, (const char*) &timeout, sizeof (timeout))) {
          LOG_WARNING("Setting TCP_USER_TIMEOUT failed! Timeout will not be set!");
        }
        // This is required for client to timeout if server lost the connection - and if the client doesn't send any data (query or heartbeat)
        timeout = timeout_milli / 1000.0;
        if (setsockopt(socket_.native(), SOL_TCP, TCP_KEEPIDLE, (const char*) &timeout, sizeof (timeout))) {
          LOG_WARNING("Setting TCP_KEEPIDLE failed! Timeout will not be set!");
        }
        timeout = 1;
        if (setsockopt(socket_.native(), SOL_TCP, TCP_KEEPINTVL, (const char*) &timeout, sizeof (timeout))) {
          LOG_WARNING("Setting TCP_KEEPINTVL failed! Timeout will not be set!");
        }
      #endif
    }

   /**
    * @brief Starts the session by connecting and listening the socket
    * Also sends the currently used IPC version.
    */
    void start() {
      write(IPCVersionMsg());
      handleConnect();
      readHeader();
    }

  /**
   * @brief Gets the IP address of the remote endpoint
   * @note the remote endpoint must be connected, otherwise empty string is returned
   * @return IP address of the remote endpoint in dot-decimal notation
   */
  std::string getRemoteIP() {
    std::error_code ec;
    std::string address = "";
    auto endpoint = socket_.remote_endpoint(ec);
    if (!ec) {
      address = endpoint.address().to_string();
    } else {
      LOG_WARNING("Cannot get remote endpoint IP address! - Error = %s", ec.message().c_str());
    }
    return address;
  }

  /**
   * @brief Gets the port of the remote endpoint
   * @note the remote endpoint must be connected, otherwise port 0 is returned
   * @return port of the remote endpoint
   */
  unsigned short getRemotePort() {
    std::error_code ec;
    unsigned short port = 0;
    auto endpoint = socket_.remote_endpoint(ec);
    if (!ec) {
      port = endpoint.port();
    } else {
      LOG_WARNING("Cannot get remote endpoint port! - Error = %s", ec.message().c_str());
    }
    return port;
  }

   /**
    * @brief Send a message through the socket
    * This is an async function and it returns immediately.
    * @param buff Buffer to send
    */
    void write(const Buffer &buff) {
      if (buff.type() == MsgType::Error) {
        return;
      }
      std::lock_guard<std::mutex> lock(out_mutex_);
      outbox_.push_back(buff);
      if (outbox_.size() > 1) {
        // We are already waiting on async_write
        return;
      }
      writeMessage();
    }

   /**
    * @brief Broadcast a message through the socket - does anything only for server sessions
    * This is an async function and it returns immediately. It will call write() deeper down as well.
    * NOTE(harijs): If both sides (Server and Client sessions) call broadcast then this can cause a cascading failure.
    * This can happen with just regular write() though as well.
    * @param buff Buffer to broadcast
    */
    inline void broadcast(const Buffer &buff) const {
      session_holder_->sessionBroadcast(buff);
    }

   /**
    * @brief Returns if sockets is subscribe to broadcasted messages
    */
    inline bool subscribed() const { return subscribed_; }

   /**
    * @brief Subscribe to broadcasted messages
    * By default all sockets are not subscribed. This must be called (often in handleConnect()) before any broadcasted messages can be recieved.
    */
    inline void subscribe() { write(Buffer(MsgType::Subscribe)); subscribed_ = true; }

   /**
    * @brief Unsubscribe to broadcasted messages
    */
    inline void unsubscribe() { write(Buffer(MsgType::Unsubscribe)); subscribed_ = false; }

   /**
    * @brief Close the socket. This happens automatically if connection ends
    */
    inline void close() {
      std::lock_guard<std::mutex> lock(shutdown_mutex_);
      if (socket_.is_open()) {
        std::error_code ec;
        socket_.shutdown(asio::ip::tcp::socket::shutdown_both, ec);

        if (ec) {
          LOG_WARNING("Session::close() - Shutdown Error = %s! Disconnected!", ec.message().c_str());
        }

        socket_.close(ec);

        if (ec) {
          LOG_WARNING("Session::close() - Close Error = %s! Disconnected!", ec.message().c_str());
        }
        connected_ = false;
      }
    }

   /**
    * @brief Returns if the socket is connected
    */
    inline bool connected() const { return connected_; }

   /**
    * @brief Virtual function that is called whenever a non-system message is recieved
    */
    virtual void handleMessage(const Buffer &message) = 0; //This is implemented for each derived session
    
  /**
    * @brief Virtual function that is called whenever a connection is established
    */
    virtual void handleConnect() {}

   /**
    * @brief Virtual function that is called whenever a disconnect happens
    */
    virtual void handleDisconnect() {}

   /**
    * @brief Returns true if session is subscribed to a certain message.
    */
    bool subscribedToMessage(const MsgType &msg_type) const {
      // Check if we don't have a filter, then by default we accept everything.
      if (subscribed_msgs_.empty()) {
        return true;
      }

      auto it = std::lower_bound(subscribed_msgs_.begin(), subscribed_msgs_.end(), msg_type);
      if (it != subscribed_msgs_.end() && *it == msg_type) {
        return true;
      }
      return false;
    }

  private:
   /**
    * @brief Sends the message via the socket if there are messages to be sent.
    * This is async and returns immediately. Works recursively until all outbound messages are sent.
    * Only one is called per socket, but is thread safe because this is called in write() function.
    */
    void writeMessage() {
      const Buffer &message = outbox_[0];
      auto self(shared_from_this());
      asio::async_write(socket_, asio::buffer(message.data(), message.length()),
        [this, self](const std::error_code &ec, std::size_t length) {
          (void) length;
          if (!ec) {
            { //Block scope the mutex
              std::lock_guard<std::mutex> lock(out_mutex_);
              outbox_.pop_front();
              if (!outbox_.empty()) { //Send more
                writeMessage();
              }
            }
          } else {
            // NOTE(harijs): These are no longer checked, because any error is considered a disconnection error. 
            // If that is not the case then we must determine which ones are correct.
            // ((ec == asio::error::eof) || (ec == asio::error::connection_reset) || 
            // (ec == asio::error::connection_aborted) || (ec == asio::error::timed_out) || (ec == asio::error::host_unreachable))
            if (connected_ == true) { // It is possible the connection was already closed in a read() function
              LOG_WARNING("Session::write() - Error = %s! Disconnected!", ec.message().c_str());
              connected_ = false;
              handleDisconnect();
              session_holder_->resolve();

              // This removes the session from server list (basically deletes itself)
              // It is safe because this lambda holds a shared_ptr as well (self)
              // For clients it doesn't do anything (as clients reconnect)
              session_holder_->sessionDisconnected();
            }
          }
        });
    }

   /**
    * @brief Reads the header when a message is recieved.
    * This is async and returns immediately.
    */
    void readHeader() {
      auto self(shared_from_this());
      asio::async_read(socket_, asio::buffer(recieve_buffer_.data(), recieve_buffer_.headerLength()),
          [this, self](const std::error_code &ec, std::size_t length) {
            (void) length;
            if (!ec) {
              readBody();
            } else {
              // NOTE(harijs): These are no longer checked, because any error is considered a disconnection error. 
              // If that is not the case then we must determine which ones are correct.
              // ((ec == asio::error::eof) || (ec == asio::error::connection_reset) || 
              // (ec == asio::error::connection_aborted) || (ec == asio::error::timed_out) || (ec == asio::error::host_unreachable))
              if (connected_ == true) { // It is possible the connection was already closed in a write() function
                LOG_WARNING("Session::read_header() - Error = %s! Disconnected!", ec.message().c_str());
                connected_ = false;
                handleDisconnect();
                session_holder_->resolve();

                // This removes the session from server list (basically deletes itself)
                // It is safe because this lambda holds a shared_ptr as well (self)
                // For clients it doesn't do anything (as clients reconnect)
                session_holder_->sessionDisconnected();
              }
            }
          });
    }

   /**
    * @brief Reads the body when a message is recieved.
    * This is async and returns immediately. If the header specifies a body, it will make an async call to get it. Otherwise it parses the message as-is.
    * Calls readHeader() in the end, so this is the main readHeader() -> readBody() -> readHeader() loop.
    */
    void readBody() {
      if (recieve_buffer_.decodeHeader()) {
        if (recieve_buffer_.bodyLength() > 0) {
          auto self(shared_from_this());
          asio::async_read(socket_, asio::buffer(recieve_buffer_.body(), recieve_buffer_.bodyLength()),
              [this, self](const std::error_code &ec, std::size_t length) {
                (void) length;
                if (!ec) {
                  switch (recieve_buffer_.type()) {
                    case MsgType::IPCVersion : {
                      ipcVersionCheck(recieve_buffer_);
                      break;
                    }
                    case MsgType::SubscribeFilter : {
                      subscribeToMessages(recieve_buffer_);
                      break;
                    }
                    case MsgType::UnsubscribeFilter : {
                      unsubscribeToMessages(recieve_buffer_);
                      break;
                    }
                    default : {
                      handleMessage(recieve_buffer_);
                      break;
                    }
                  }
                  readHeader();
                } else if ((ec == asio::error::eof) || (ec == asio::error::connection_reset) || (ec == asio::error::connection_aborted)) {
                  if (connected_ == true) { // It is possible the connection was already closed in a write() function
                    LOG_WARNING("Session::read_body() - Error = %s! Reconnection will be attempted!", ec.message().c_str());
                    connected_ = false;
                    handleDisconnect();
                    session_holder_->resolve();

                    // This removes the session from server list (basically deletes itself)
                    // It is safe because this lambda holds a shared_ptr as well (self)
                    // For clients it doesn't do anything (as clients reconnect)
                    session_holder_->sessionDisconnected();
                  }
                } else {
                  LOG_ERROR("Session::read_body() - Error = %s", ec.message().c_str());
                }
              });
        } else { //This is a query message and doesn't have data
          switch (recieve_buffer_.type()) {
            case MsgType::Subscribe : {
              subscribed_ = true;
              break;
            }
            case MsgType::Unsubscribe : {
              subscribed_ = false;
              break;
            }
            case MsgType::ClearSubscriptionFilter : {
              clearSubscriptionToMessages();
              break;
            }
            default : {
              handleMessage(recieve_buffer_);
              break;
            }
          }
          readHeader();
        }
      }
    }

   /**
    * @brief Adds message types to the subscription filter
    * If the session isn't subscribed, then it is after this call.
    * Inserts a sorted vector inside a sorted vector. So both vectors must be always sorted.
    * Optimizes for long inserts at the end (O(m)).
    * Each binary search starts at the end of the previous search (at most m*log2(n)).
    */
    void subscribeToMessages(const Buffer &message) {
      auto msg = message.cast_array<ipc::SubscribeFilterMsg>();
      size_t start_distance = 0;
      size_t push_from = 0;
      bool pushing_end = false;
      for (const auto &m : msg) {
        //TODO(harijs): Make a custom std::lower_bound function that can bail out early
        auto start = subscribed_msgs_.begin();
        std::advance(start, start_distance);
        // If we are at the end, from now on we can just copy to end, so break this loop
        if (start == subscribed_msgs_.end()) {
          //NOTE(harijs): This will break if we us non-continous data structure! So use vector!
          push_from = std::distance(&*msg.begin(), &m);
          pushing_end = true;
          break;
        } else {
          auto it = std::lower_bound(start, subscribed_msgs_.end(), m.msg_type);
          if (it != subscribed_msgs_.end() && *it == m.msg_type) { //Skip messages if we already have them
            continue;
          }
          start_distance = std::distance(subscribed_msgs_.begin(), it)+1;
          subscribed_msgs_.insert(it, m.msg_type);
        }
      }

      // We copy the rest to the end if necessary
      if (pushing_end == true) {
        //TODO(harijs): std::insert() might be faster, but would require a custom inserter
        // std::copy with reserve is also faster, but we cannot do it with SubscribeFilterMsg and MsgType (they are different types)
        size_t current_size = subscribed_msgs_.size();
        subscribed_msgs_.resize(current_size + (msg.size() - push_from));
        push_from -= current_size;
        for (size_t i = current_size; i < subscribed_msgs_.size(); ++i) {
          subscribed_msgs_[i] = msg[push_from + i].msg_type;
        }
        // This is slower
        //std::transform(msg.begin(), msg.end(), std::back_inserter(subscribed_msgs_),
        //       [](const ipc::SubscribeFilterMsg& msg) { return msg.msg_type; });
      }

      subscribed_ = true;
    }

   /**
    * @brief Removes message types from the subscription filter
    * Session will stay subscribed even if the filter is empty after this function call!
    * Removes a sorted vector from a sorted vector. So both vectors must be always sorted.
    * Optimizes for long removes at any point (O(m)).
    * Breaks at first missing item past the end (because the rest won't be there either).
    */
    void unsubscribeToMessages(const Buffer &message) {
      auto msg = message.cast_array<ipc::UnsubscribeFilterMsg>();
      size_t start_distance = 0;
      bool on_roll = false;
      auto roll_start = subscribed_msgs_.end();
      auto roll_end = subscribed_msgs_.end();
      for (const auto &m : msg) {
        //TODO(harijs): Make a custom std::lower_bound function that can bail out early
        auto start = subscribed_msgs_.begin();
        std::advance(start, start_distance);

        // If we are at the end, we can bail out - nothing can be found past here
        if (start == subscribed_msgs_.end()) {
          break;
        } else {
          // We are on a roll! The next iteration hit the next target.
          if (*start == m.msg_type) {
            if (on_roll == false) {
              on_roll = true;
              roll_start = start;
              roll_end = start;
            }
            roll_end++;
            start_distance++;
          } else {
            auto it = std::lower_bound(start, subscribed_msgs_.end(), m.msg_type);
            if (it != subscribed_msgs_.end() && *it == m.msg_type) {
              // Roll ended - delete everything we found
              if (on_roll == true) {
                it -= std::distance(roll_start, roll_end);
                subscribed_msgs_.erase(roll_start, roll_end);
                on_roll = false;
              }
              // We cannot skip this by erasing to roll_end+1, because "roll_end+1 != it" in most cases
              it = subscribed_msgs_.erase(it);
            }
            start_distance = std::distance(subscribed_msgs_.begin(), it);
          }
        }
      }

      // In case we were on a roll till the end
      if (on_roll == true) {
        if (roll_start == subscribed_msgs_.begin() && roll_end == subscribed_msgs_.end()) {
          subscribed_msgs_.clear();
        } else {
          subscribed_msgs_.erase(roll_start, roll_end);
        }
      }
    }

   /**
    * @brief Removes all message types from the subscription filters by clearing it
    */
    void clearSubscriptionToMessages() {
      subscribed_msgs_.clear();
    }

   /**
    * @brief Checks versions - if major mismatch then disconnect - if minor mismatch then a warning
    */
    void ipcVersionCheck(const Buffer &message) {
      auto msg = message.cast<ipc::IPCVersionMsg>();
      if (msg->major != ipc::major_version) {
        LOG_ERROR("IPC Major Version mismatch - local %u.%u - remote %u.%u", 
                  static_cast<unsigned>(ipc::major_version), static_cast<unsigned>(ipc::minor_version), 
                  static_cast<unsigned>(msg->major), static_cast<unsigned>(msg->minor));
        handleDisconnect();
        close();
        session_holder_->sessionDisconnected();
      } else if (msg->minor != ipc::minor_version) {
        LOG_WARNING("IPC Minor Version mismatch - local %u.%u - remote %u.%u", 
                  static_cast<unsigned>(ipc::major_version), static_cast<unsigned>(ipc::minor_version), 
                  static_cast<unsigned>(msg->major), static_cast<unsigned>(msg->minor));
      }
    }

  private:
    asio::ip::tcp::socket socket_;
    std::deque<Buffer> outbox_;

    // We might not need this if we don't call handle_message in a different thread, so recieve_buffer_ is used for now
    //std::deque<Buffer> inbox_;

    std::mutex out_mutex_;
    //std::mutex in_mutex_;

    std::mutex shutdown_mutex_;

    Buffer recieve_buffer_;
    
    SessionHolder* session_holder_; // We don't use a smart pointer here because nothing is owned

    bool subscribed_ = false;
    bool connected_ = true;
    bool client_ = false; // True for client session, false for server session

    std::vector<MsgType> subscribed_msgs_; // List of all messages we want to recieve via broadcast()
};

}

#endif //IPC_SESSION_H