#pragma once

#include <asio.hpp>
#include <cstdint>
#include <map>
#include <memory>
#include <mutex>
#include <sstream>
#include <string>
#include <utility>
#include <vector>

// Minimal Prometheus exporter for the secs-gem runtime.
//
// Design:
//   * Registry: counters + gauges keyed by (name, labels).
//   * PrometheusServer: asio acceptor; every HTTP request gets the
//     same /metrics text body.  No path routing, no auth — drop a
//     reverse proxy in front for either.
//   * StandardMetrics: pre-wired hooks for Connection + EquipmentDataModel
//     (RX/TX counters, selected/closed gauges, FSM transitions).
//
// Thread-safety: the Registry uses a mutex internally so updates from
// the io_context thread and scrapes from the io's HTTP handler are
// safe in any order.  Lock contention is negligible at typical fab
// rates (a few hundred mutations/sec).
namespace secsgem::metrics {

using Labels = std::vector<std::pair<std::string, std::string>>;

enum class MetricType : uint8_t { Counter, Gauge };

class Registry {
 public:
  // Optional: register a HELP line and type declaration.  Scrapers
  // tolerate omission, but Grafana / alertmanager queries become
  // self-documenting with them.
  void describe(const std::string& name, const std::string& help,
                MetricType type) {
    std::lock_guard<std::mutex> g(mu_);
    descriptions_[name] = {help, type};
  }

  void inc(const std::string& name, const Labels& labels = {},
           double delta = 1.0) {
    std::lock_guard<std::mutex> g(mu_);
    counters_[Key{name, labels}] += delta;
  }

  void set_gauge(const std::string& name, double value,
                 const Labels& labels = {}) {
    std::lock_guard<std::mutex> g(mu_);
    gauges_[Key{name, labels}] = value;
  }

  // Render in Prometheus text-exposition format.
  std::string render() const {
    std::lock_guard<std::mutex> g(mu_);
    std::ostringstream os;

    // Group by name so each # HELP / # TYPE pair appears once.
    std::map<std::string, std::vector<std::pair<Labels, double>>> by_name;
    for (const auto& [k, v] : counters_) by_name[k.name].push_back({k.labels, v});

    auto emit_header = [&](const std::string& name, MetricType inferred) {
      auto it = descriptions_.find(name);
      if (it != descriptions_.end()) {
        os << "# HELP " << name << " " << it->second.help << "\n";
        os << "# TYPE " << name
           << (it->second.type == MetricType::Counter ? " counter" : " gauge")
           << "\n";
      } else {
        os << "# TYPE " << name
           << (inferred == MetricType::Counter ? " counter" : " gauge")
           << "\n";
      }
    };

    for (const auto& [name, rows] : by_name) {
      emit_header(name, MetricType::Counter);
      for (const auto& [labels, v] : rows) {
        os << name;
        emit_labels(os, labels);
        os << " " << v << "\n";
      }
    }

    by_name.clear();
    for (const auto& [k, v] : gauges_) by_name[k.name].push_back({k.labels, v});
    for (const auto& [name, rows] : by_name) {
      emit_header(name, MetricType::Gauge);
      for (const auto& [labels, v] : rows) {
        os << name;
        emit_labels(os, labels);
        os << " " << v << "\n";
      }
    }
    return os.str();
  }

 private:
  struct Key {
    std::string name;
    Labels labels;
    bool operator<(const Key& other) const {
      if (name != other.name) return name < other.name;
      return labels < other.labels;
    }
  };

  static void emit_labels(std::ostream& os, const Labels& labels) {
    if (labels.empty()) return;
    os << "{";
    for (std::size_t i = 0; i < labels.size(); ++i) {
      if (i) os << ",";
      os << labels[i].first << "=\"" << labels[i].second << "\"";
    }
    os << "}";
  }

  struct Description {
    std::string help;
    MetricType type;
  };

  mutable std::mutex mu_;
  std::map<Key, double> counters_;
  std::map<Key, double> gauges_;
  std::map<std::string, Description> descriptions_;
};

// Tiny HTTP server: accepts on `port`, replies to any GET with the
// rendered registry.  No keep-alive, no path routing — drop nginx or
// Caddy in front for those.
class PrometheusServer : public std::enable_shared_from_this<PrometheusServer> {
 public:
  PrometheusServer(asio::io_context& io, uint16_t port,
                   std::shared_ptr<Registry> reg)
      : acceptor_(io, asio::ip::tcp::endpoint(asio::ip::tcp::v4(), port)),
        registry_(std::move(reg)) {}

  void start() { do_accept_(); }

  uint16_t port() const {
    return acceptor_.local_endpoint().port();
  }

 private:
  void do_accept_() {
    auto self = shared_from_this();
    acceptor_.async_accept([self](std::error_code ec,
                                   asio::ip::tcp::socket sock) {
      if (!ec) self->serve_(std::move(sock));
      if (self->acceptor_.is_open()) self->do_accept_();
    });
  }

  void serve_(asio::ip::tcp::socket sock) {
    auto socket = std::make_shared<asio::ip::tcp::socket>(std::move(sock));
    auto buf = std::make_shared<asio::streambuf>();
    auto self = shared_from_this();
    asio::async_read_until(
        *socket, *buf, "\r\n\r\n",
        [self, socket, buf](std::error_code ec, std::size_t) {
          if (ec) return;
          const auto body = self->registry_->render();
          auto resp = std::make_shared<std::string>();
          *resp = "HTTP/1.1 200 OK\r\n"
                  "Content-Type: text/plain; version=0.0.4\r\n"
                  "Content-Length: " + std::to_string(body.size()) + "\r\n"
                  "Connection: close\r\n\r\n" + body;
          asio::async_write(*socket, asio::buffer(*resp),
                            [resp, socket](std::error_code, std::size_t) {
                              std::error_code shut;
                              socket->shutdown(
                                  asio::ip::tcp::socket::shutdown_both, shut);
                            });
        });
  }

  asio::ip::tcp::acceptor acceptor_;
  std::shared_ptr<Registry> registry_;
};

}  // namespace secsgem::metrics