/* -*- c++ -*- */
/* 
 * Copyright 2021 Satlab A/S.
 * 
 * This is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3, or (at your option)
 * any later version.
 * 
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this software; see the file COPYING.  If not, write to
 * the Free Software Foundation, Inc., 51 Franklin Street,
 * Boston, MA 02110-1301, USA.
 */

#include <websocketpp/config/asio_no_tls.hpp>
#include <websocketpp/common/thread.hpp>
#include <websocketpp/server.hpp>
#include <websocketpp/client.hpp>

#include <iostream>
#include <set>

#include "broadcast_server.h"

typedef websocketpp::server<websocketpp::config::asio> server;

using websocketpp::connection_hdl;
using websocketpp::lib::placeholders::_1;
using websocketpp::lib::bind;

using websocketpp::lib::thread;
using websocketpp::lib::lock_guard;

namespace gr {
  namespace satlab {

    broadcast_server::broadcast_server() {
      // Initialize Asio Transport
      m_server.init_asio();

      // No logging
      m_server.clear_access_channels(websocketpp::log::alevel::all);
      m_server.clear_error_channels(websocketpp::log::alevel::all);

      // Allow address reuse
      m_server.set_reuse_addr(true);

      // Register handler callbacks
      m_server.set_open_handler(bind(&broadcast_server::on_open,this,_1));
      m_server.set_close_handler(bind(&broadcast_server::on_close,this,_1));
    }

    void broadcast_server::run(uint16_t port) {
      // listen on specified port
      m_server.listen(port);

      // Start the server accept loop
      m_server.start_accept();

      // Start the ASIO io_service run loop
      m_thread_server = new websocketpp::lib::thread(bind(&server::run, &m_server));
    }

    void broadcast_server::stop() {
      // Stop server accept loop
      m_server.stop_listening();

      // Close all established connections
      for (auto it : m_connections) {
        try {
          websocketpp::lib::error_code ec;
          m_server.close(it, websocketpp::close::status::going_away, "", ec);
        } catch (const std::exception & e) {
          std::cout << e.what() << std::endl;
        }
      }

      // Wait for ASIO server to exit
      m_thread_server->join();
    }

    void broadcast_server::on_open(connection_hdl hdl) {
      std::lock_guard<std::mutex> lock(m_mutex);
      try {
	      m_server.send(hdl, m_greeting, websocketpp::frame::opcode::text);
      } catch (const std::exception & e) {
        std::cout << e.what() << std::endl;
        return;
      }
      m_connections.insert(hdl);
    }

    void broadcast_server::on_close(connection_hdl hdl) {
      std::lock_guard<std::mutex> lock(m_mutex);
      m_connections.erase(hdl);
    }

    void broadcast_server::broadcast(std::string msg) {
      std::lock_guard<std::mutex> lock(m_mutex);    
      for (auto it : m_connections) {
        try {
          // Don't bother pushing data if client can't keep up
          websocketpp::client<websocketpp::config::asio>::connection_ptr connection = m_server.get_con_from_hdl(it);
          if (connection->get_buffered_amount() < 1000000)
            m_server.send(it, msg, websocketpp::frame::opcode::text);
        } catch (const std::exception & e) {
          std::cout << e.what() << std::endl;
          m_connections.erase(it);
        }
      }
    }

    void broadcast_server::set_greeting(std::string msg) {
      std::lock_guard<std::mutex> lock(m_mutex);
      m_greeting = msg;
      for (auto it : m_connections) {
        try {
          m_server.send(it, m_greeting, websocketpp::frame::opcode::text);
        } catch (const std::exception & e) {
          std::cout << e.what() << std::endl;
          m_connections.erase(it);
        }
      }
    }

  }
}
