#pragma once

#include <cstdint>
#include <functional>
#include <map>
#include <optional>
#include <utility>

#include "secsgem/secs2/message.hpp"

namespace secsgem::gem {

namespace s2 = secsgem::secs2;

// A small (stream, function) dispatch table. The Server registers one
// handler per primary SxFy and calls `dispatch` from the Connection's
// message handler. Replaces the imperative if-ladder; behaviour stays in
// the handlers (since each SxFy reply shape is unique), but routing is
// data.
//
// Default behaviour for unregistered primaries:
//   - If a `fallback` is installed, it runs.
//   - Otherwise, if the inbound message has W set, reply with SxF0
//     (Abort) per E5 convention.
//   - Otherwise, do nothing.
class Router {
 public:
  using Handler = std::function<std::optional<s2::Message>(const s2::Message&)>;

  Router& on(uint8_t stream, uint8_t function, Handler h) {
    handlers_[{stream, function}] = std::move(h);
    return *this;
  }

  Router& fallback(Handler h) {
    fallback_ = std::move(h);
    return *this;
  }

  std::optional<s2::Message> dispatch(const s2::Message& msg) const {
    auto it = handlers_.find({msg.stream, msg.function});
    if (it != handlers_.end()) return it->second(msg);
    if (fallback_) return fallback_(msg);
    if (msg.reply_expected) return s2::Message(msg.stream, 0, false);
    return std::nullopt;
  }

  std::size_t size() const { return handlers_.size(); }

 private:
  std::map<std::pair<uint8_t, uint8_t>, Handler> handlers_;
  Handler fallback_;
};

}  // namespace secsgem::gem