#include "connection.hpp"

#include "linebuffer.hpp"

#include <mybase64.hpp>

#include <boost/log/trivial.hpp>

namespace {
#include "irc_commands.inc"
} // namespace

using namespace std::literals;

Connection::Connection(boost::asio::io_context &io)
    : stream_{io}
    , watchdog_timer_{io}
    , write_posted_{false}
    , stalled_{false}
{
}

auto Connection::write_buffers() -> void
{
    std::vector<boost::asio::const_buffer> buffers;
    buffers.reserve(write_strings_.size());
    for (const auto &elt : write_strings_)
    {
        buffers.push_back(boost::asio::buffer(elt));
    }
    boost::asio::async_write(
        stream_,
        buffers,
        [this, strings = std::move(write_strings_)](const boost::system::error_code &error, std::size_t) {
            if (not error)
            {
                if (write_strings_.empty())
                {
                    write_posted_ = false;
                }
                else
                {
                    write_buffers();
                }
            }
        }
    );
    write_strings_.clear();
}

auto Connection::connect(
    boost::asio::io_context &io,
    std::string host,
    std::string port
) -> boost::asio::awaitable<void>
{
    using namespace std::placeholders;

    // keep connection alive while coroutine is active
    const auto self = shared_from_this();

    {
        auto resolver = boost::asio::ip::tcp::resolver{io};
        const auto endpoints = co_await resolver.async_resolve(host, port, boost::asio::use_awaitable);
        const auto endpoint = co_await boost::asio::async_connect(stream_, endpoints, boost::asio::use_awaitable);

        BOOST_LOG_TRIVIAL(debug) << "CONNECTED: " << endpoint;
        sig_connect();
    }

    watchdog();

    for (LineBuffer buffer{32'768};;)
    {
        boost::system::error_code error;
        const auto n = co_await stream_.async_read_some(buffer.get_buffer(), boost::asio::redirect_error(boost::asio::use_awaitable, error));
        if (error)
        {
            break;
        }
        buffer.add_bytes(n, [this](char *line) {
            BOOST_LOG_TRIVIAL(debug) << "RECV: " << line;
            watchdog_activity();
            dispatch_line(line);
        });
    }

    watchdog_timer_.cancel();
    stream_.close();

    BOOST_LOG_TRIVIAL(debug) << "DISCONNECTED";
    sig_disconnect();
}

auto Connection::watchdog() -> void
{
    watchdog_timer_.expires_after(watchdog_duration);
    watchdog_timer_.async_wait([this](const auto &error) {
        if (not error)
        {
            if (stalled_)
            {
                BOOST_LOG_TRIVIAL(debug) << "Watchdog timer elapsed, closing stream";
                close();
            }
            else
            {
                send_ping("watchdog");
                stalled_ = true;
                watchdog();
            }
        }
    });
}

auto Connection::watchdog_activity() -> void
{
    stalled_ = false;
    watchdog_timer_.expires_after(watchdog_duration);
}

/// Parse IRC message line and dispatch it to the ircmsg slot.
auto Connection::dispatch_line(char *line) -> void
{
    const auto msg = parse_irc_message(line);
    const auto recognized = IrcCommandHash::in_word_set(msg.command.data(), msg.command.size());
    const auto command
        = recognized
            && recognized->min_args <= msg.args.size()
            && recognized->max_args >= msg.args.size()
        ? recognized->command
        : IrcCommand::UNKNOWN;

    switch (command)
    {

    // Respond to pings immediate and discard
    case IrcCommand::PING:
        send_pong(msg.args[0]);
        break;

    // Unknown message generate warnings but do not dispatch
    // Messages can be unknown due to bad command or bad argument count
    case IrcCommand::UNKNOWN:
        BOOST_LOG_TRIVIAL(warning) << "Unrecognized command: " << msg.command << " " << msg.args.size();
        break;

    case IrcCommand::AUTHENTICATE:
        on_authenticate(msg.args[0]);
        break;

    // Server notice generate snote events but not IRC command events
    case IrcCommand::NOTICE:
        if (auto match = snoteCore.match(msg))
        {
            sig_snote(*match);
            break;
        }
        /* FALLTHROUGH */

    // Normal IRC commands
    default:
        sig_ircmsg(command, msg);
        break;
    }
}

auto Connection::write_line(std::string message) -> void
{
    BOOST_LOG_TRIVIAL(debug) << "SEND: " << message;
    message += "\r\n";
    write_strings_.push_back(std::move(message));

    if (not write_posted_)
    {
        write_posted_ = true;
        boost::asio::post(stream_.get_executor(), [weak = weak_from_this()]() {
            if (auto self = weak.lock())
            {
                self->write_buffers();
            }
        });
    }
}

auto Connection::close() -> void
{
    stream_.close();
}

static auto is_invalid_last(char x) -> bool
{
    return x == '\0' || x == '\r' || x == '\n';
}

auto Connection::write_irc(std::string message) -> void
{
    write_line(std::move(message));
}

auto Connection::write_irc(std::string front, std::string_view last) -> void
{
    if (last.end() != std::find_if(last.begin(), last.end(), is_invalid_last))
    {
        throw std::runtime_error{"bad irc argument"};
    }

    front += " :";
    front += last;
    write_irc(std::move(front));
}

auto Connection::send_ping(std::string_view txt) -> void
{
    write_irc("PING", txt);
}

auto Connection::send_pong(std::string_view txt) -> void
{
    write_irc("PONG", txt);
}

auto Connection::send_pass(std::string_view password) -> void
{
    write_irc("PASS", password);
}

auto Connection::send_user(std::string_view user, std::string_view real) -> void
{
    write_irc("USER", user, "*", "*", real);
}

auto Connection::send_nick(std::string_view nick) -> void
{
    write_irc("NICK", nick);
}

auto Connection::send_cap_ls() -> void
{
    write_irc("CAP", "LS", "302");
}

auto Connection::send_cap_end() -> void
{
    write_irc("CAP", "END");
}

auto Connection::send_cap_req(std::string_view caps) -> void
{
    write_irc("CAP", "REQ", caps);
}

auto Connection::send_privmsg(std::string_view target, std::string_view message) -> void
{
    write_irc("PRIVMSG", target, message);
}

auto Connection::send_notice(std::string_view target, std::string_view message) -> void
{
    write_irc("NOTICE", target, message);
}

auto Connection::send_authenticate(std::string_view message) -> void
{
    write_irc("AUTHENTICATE", message);
}

auto Connection::on_authenticate(const std::string_view chunk) -> void
{
    if (chunk != "+"sv)
    {
        authenticate_buffer_ += chunk;
    }

    if (chunk.size() != 400)
    {
        std::string decoded;
        decoded.resize(mybase64::decoded_size(authenticate_buffer_.size()));
        std::size_t len;

        if (mybase64::decode(authenticate_buffer_, decoded.data(), &len))
        {
            decoded.resize(len);
            sig_authenticate(decoded);
        }
        else
        {
            BOOST_LOG_TRIVIAL(debug) << "Invalid AUTHENTICATE base64"sv;
            send_authenticate("*"sv); // abort SASL
        }

        authenticate_buffer_.clear();
    }
    else if (authenticate_buffer_.size() > 1024)
    {
        BOOST_LOG_TRIVIAL(debug) << "AUTHENTICATE buffer overflow"sv;
        authenticate_buffer_.clear();
        send_authenticate("*"sv); // abort SASL
    }
}

auto Connection::send_authenticate_abort() -> void
{
    send_authenticate("*");
}

auto Connection::send_authenticate_encoded(std::string_view body) -> void
{
    std::string encoded(mybase64::encoded_size(body.size()), 0);
    mybase64::encode(body, encoded.data());

    for (size_t lo = 0; lo < encoded.size(); lo += 400)
    {
        const auto hi = std::min(lo + 400, encoded.size());
        const std::string_view chunk{encoded.begin() + lo, encoded.begin() + hi};
        send_authenticate(chunk);
    }

    if (encoded.size() % 400 == 0)
    {
        send_authenticate("+"sv);
    }
}