#include "port_scanner.hpp"
#include "logger.hpp"
#include <boost/asio.hpp>
#include <string>

PortScanner::PortScanner(
    const std::string &ip, const std::vector<uint16_t> &ports
)
    : ports_(ports)
{
    ErrorCode ec;
    addr_ = boost::asio::ip::make_address(ip, ec);
    ip_ok_ = !ec;

    if (!ip_ok_) {
        Logger::error(std::string("PortScanner: invalid IP: ") + ip);
    } else {
        Logger::debug(std::string("PortScanner: using IP: ") + ip);
    }
}

void PortScanner::setTimeoutMs(int ms)
{
    if (ms > 0) {
        timeout_ms_ = ms;
        Logger::debug(
            std::string("PortScanner: timeout set to ")
            + std::to_string(timeout_ms_) + " ms"
        );
    }
}

void PortScanner::run()
{
    if (!ip_ok_) {
        Logger::error("PortScanner::run - aborting: ip not ok");
        return;
    }
    for (uint16_t p : ports_) {
        if (probe(p)) {
            Logger::info(std::to_string(p));
        }
    }
}

bool PortScanner::probe(uint16_t port)
{
    TcpSocket s(io_);
    SteadyTimer t(io_);
    auto ep = endpoint(port);

    Logger::debug(
        std::string("PortScanner::probe - probing port ") + std::to_string(port)
    );

    armTimer(t, s);
    armConnect(s, ep);
    runUntilDone();
    cancelTimer(t);
    drainLeftovers();
    close(s);

    if (ok_) {
        Logger::debug(std::string("Port ") + std::to_string(port) + " is OPEN");
    } else {
        Logger::debug(
            std::string("Port ") + std::to_string(port) + " is CLOSED/FILTERED"
        );
    }

    return ok_;
}

TcpEndpoint PortScanner::endpoint(uint16_t port) const
{
    return TcpEndpoint(addr_, port);
}

void PortScanner::armTimer(SteadyTimer &t, TcpSocket &s)
{
    t.expires_after(std::chrono::milliseconds(timeout_ms_));

    auto onTimeout = [&s](const ErrorCode &ec) {
        if (!ec) {
            Logger::debug("Timer expired -> cancelling socket");
            s.cancel();
        }
    };

    t.async_wait(onTimeout);
}

void PortScanner::armConnect(TcpSocket &s, const TcpEndpoint &ep)
{
    ok_ = false;
    done_ = false;

    auto onConnect = [this](const ErrorCode &ec) {
        if (!ec) {
            ok_ = true;
            Logger::debug("async_connect succeeded");
        } else {
            ok_ = false;
            Logger::debug(std::string("async_connect failed: ") + ec.message());
        }
        done_ = true;
    };

    s.async_connect(ep, onConnect);
}

void PortScanner::runUntilDone()
{
    io_.restart();
    while (!done_ && io_.run_one()) {
    }
}

void PortScanner::cancelTimer(SteadyTimer &t)
{
    ErrorCode ec;
    t.cancel();
    if (ec) {
        Logger::warn(std::string("cancelTimer: ") + ec.message());
    }
}

void PortScanner::drainLeftovers()
{
    io_.restart();
    while (io_.run_one()) {
        continue;
    }
}

void PortScanner::close(TcpSocket &s)
{
    ErrorCode ig;
    if (s.close(ig)) {
        Logger::warn(std::string("close socket failed: ") + ig.message());
    }
}