#include "socks5/server.h"

#include "common/error.h"
#include "common/logging.h"
#include "common/metrics.h"

namespace ovg::socks5 {
namespace {

constexpr const char *kMod = "socks5";

metrics::Counter *m_accepted() {
  static auto *c = metrics::counter("ovg_socks5_sessions_total",
                                    "Accepted SOCKS5 connections");
  return c;
}
metrics::Counter *m_rejected() {
  static auto *c =
      metrics::counter("ovg_socks5_rejected_total",
                       "Connections refused by max_sessions admission control");
  return c;
}
metrics::Counter *m_no_egress() {
  static auto *c = metrics::counter("ovg_socks5_no_egress_total",
                                    "Connections refused: no usable tunnel");
  return c;
}
metrics::Gauge *g_active() {
  static auto *g =
      metrics::gauge("ovg_socks5_sessions_active", "Live SOCKS5 sessions");
  return g;
}

}  // namespace

Server::Server(asio::io_context &io, const Config &cfg, EgressProvider acquire)
    : io_(io),
      cfg_(cfg.socks5),
      acquire_(std::move(acquire)),
      accept_strand_(asio::make_strand(io)),
      acceptor_(accept_strand_),
      auth_(cfg.socks5.users, cfg.socks5.require_auth) {
  auto opts = std::make_shared<SessionOptions>();
  opts->socks5 = cfg.socks5;
  opts->retry_zero_progress = cfg.switching.retry_zero_progress;
  opts->rehome_udp = cfg.switching.rehome_udp;
  opts_ = std::move(opts);
}

Server::~Server() { stop(); }

bool Server::start(std::string *err) {
  std::error_code ec;
  const auto addr = asio::ip::make_address(cfg_.listen_address, ec);
  if (ec) {
    *err = "socks5.listen_address is not an IP address: " + cfg_.listen_address;
    return false;
  }
  const asio::ip::tcp::endpoint ep(addr, cfg_.listen_port);

  acceptor_.open(ep.protocol(), ec);
  if (ec) {
    *err = "cannot open listening socket: " + ec.message();
    return false;
  }
  acceptor_.set_option(asio::socket_base::reuse_address(true), ec);
  acceptor_.bind(ep, ec);
  if (ec) {
    *err = "cannot bind " + ep.address().to_string() + ":" +
           std::to_string(ep.port()) + ": " + ec.message();
    return false;
  }
  // A deep backlog matters here: 1000 clients reconnecting after a switch
  // arrive in a burst, and the default of 5 would turn that into connection
  // refused.
  acceptor_.listen(asio::socket_base::max_listen_connections, ec);
  if (ec) {
    *err = "cannot listen: " + ec.message();
    return false;
  }
  port_ = acceptor_.local_endpoint(ec).port();

  LOG_INFO(kMod, "listening on {}:{} (auth {}, max_sessions {}, udp {})",
           ep.address().to_string(), port_,
           cfg_.require_auth ? "required" : "optional", cfg_.max_sessions,
           cfg_.udp_associate_enabled ? "enabled" : "disabled");

  do_accept();
  return true;
}

void Server::do_accept() {
  if (stopping_.load(std::memory_order_acquire)) return;

  // Accepting straight onto a fresh strand gives the session a socket whose
  // executor already is its strand -- see socks5/session.h.
  acceptor_.async_accept(
      asio::make_strand(io_),
      [this](const std::error_code &ec, ClientSocket sock) {
        if (ec) {
          if (ec == asio::error::operation_aborted) return;
          LOG_WARN(kMod, "accept failed: {}", ec.message());
          // An accept error is usually per-connection (EMFILE, ECONNABORTED).
          // Give up the listener only if it was closed under us.
          if (acceptor_.is_open()) do_accept();
          return;
        }

        if (stopping_.load(std::memory_order_acquire)) {
          std::error_code ignored;
          sock.close(ignored);
          return;
        }

        // Admission control, before anything else allocates.
        if (active_.load(std::memory_order_relaxed) >=
            static_cast<int64_t>(cfg_.max_sessions)) {
          rejected_.fetch_add(1, std::memory_order_relaxed);
          m_rejected()->inc();
          LOG_WARN(kMod, "refusing connection: {} sessions already live",
                   cfg_.max_sessions);
          std::error_code ignored;
          sock.close(ignored);
          do_accept();
          return;
        }

        egress::EgressPtr eg = acquire_ ? acquire_() : nullptr;
        if (!eg) {
          // Closing without a SOCKS5 reply is deliberate: we have not read the
          // greeting yet, so there is no negotiated framing to reply in.
          no_egress_.fetch_add(1, std::memory_order_relaxed);
          m_no_egress()->inc();
          LOG_WARN(kMod, "refusing connection: no usable egress");
          std::error_code ignored;
          sock.close(ignored);
          do_accept();
          return;
        }

        uint64_t id;
        SessionPtr s;
        {
          std::lock_guard<std::mutex> lk(mu_);
          id = next_id_++;
          s = Session::create(id, std::move(sock), opts_, &auth_, std::move(eg),
                              [this](uint64_t sid) { reap(sid); });
          sessions_.emplace(id, s);
        }
        active_.fetch_add(1, std::memory_order_relaxed);
        g_active()->add(1);
        accepted_.fetch_add(1, std::memory_order_relaxed);
        m_accepted()->inc();

        s->start();
        do_accept();
      });
}

void Server::reap(uint64_t id) {
  {
    std::lock_guard<std::mutex> lk(mu_);
    if (sessions_.erase(id) == 0) return;  // already reaped
  }
  active_.fetch_sub(1, std::memory_order_relaxed);
  g_active()->sub(1);
}

void Server::stop() {
  if (stopping_.exchange(true, std::memory_order_acq_rel)) return;

  asio::post(accept_strand_, [this] {
    std::error_code ignored;
    acceptor_.close(ignored);
  });

  for (const auto &s : snapshot()) s->force_close("server shutting down");
  LOG_INFO(kMod, "listener stopped");
}

std::vector<SessionPtr> Server::snapshot() const {
  std::vector<SessionPtr> out;
  std::lock_guard<std::mutex> lk(mu_);
  out.reserve(sessions_.size());
  for (const auto &kv : sessions_) {
    if (auto s = kv.second.lock()) out.push_back(std::move(s));
  }
  return out;
}

void Server::on_promote(const egress::EgressPtr &old_e,
                        const egress::EgressPtr &new_e) {
  if (!old_e || !new_e) return;
  // Offer, do not command: each session decides on its own strand whether it is
  // actually re-homeable. Everything that declines stays on the old egress and
  // drains normally.
  const auto live = snapshot();
  for (const auto &s : live) s->try_rehome(old_e, new_e);
  LOG_INFO(kMod, "offered {} live session(s) a move from {} to {}", live.size(),
           old_e->label(), new_e->label());
}

void Server::on_drain_expired(const egress::EgressPtr &e) {
  if (!e) return;
  size_t n = 0;
  for (const auto &s : snapshot()) {
    s->close_if_on(e, "drain grace expired");
    ++n;
  }
  LOG_WARN(kMod, "drain window for {} expired; closing whatever is left of {} "
                 "session(s)",
           e->label(), n);
}

Server::Stats Server::stats() const {
  Stats s;
  s.accepted = accepted_.load(std::memory_order_relaxed);
  s.rejected = rejected_.load(std::memory_order_relaxed);
  s.no_egress = no_egress_.load(std::memory_order_relaxed);
  s.active = active_.load(std::memory_order_relaxed);
  s.auth_ok = auth_.successes();
  s.auth_failed = auth_.failures();
  return s;
}

std::vector<Session::Info> Server::sessions(size_t limit) const {
  std::vector<Session::Info> out;
  for (const auto &s : snapshot()) {
    if (out.size() >= limit) break;
    out.push_back(s->info());
  }
  return out;
}

}  // namespace ovg::socks5