// HSMS-GS (multi-session, E37 §11) tests.
//
// The same hsms::Connection class supports HSMS-SS by default (one
// session registered by the constructor) and HSMS-GS once
// `add_session(device_id)` registers more.  These tests cover the
// distinct behaviours: per-session selected state, session-routed
// data dispatch, rejection of data on the wrong session, and
// Reject(EntityNotSelected) for a Select.req targeting an unknown
// session id.

#include <doctest/doctest.h>

#include <asio.hpp>
#include <atomic>
#include <chrono>
#include <cstdint>
#include <memory>
#include <optional>
#include <thread>
#include <utility>
#include <vector>

#include "secsgem/hsms/connection.hpp"
#include "secsgem/hsms/header.hpp"
#include "secsgem/secs2/message.hpp"

using namespace secsgem;
using namespace std::chrono_literals;

namespace {

struct SocketPair {
  asio::io_context io;
  asio::ip::tcp::socket a{io}, b{io};
  SocketPair() {
    asio::ip::tcp::acceptor acc(io, asio::ip::tcp::endpoint(
                                        asio::ip::address_v4::loopback(), 0));
    const auto port = acc.local_endpoint().port();
    bool da = false, db = false;
    std::error_code ea, eb;
    acc.async_accept(a, [&](std::error_code ec) { ea = ec; da = true; });
    b.async_connect({asio::ip::address_v4::loopback(), port},
                    [&](std::error_code ec) { eb = ec; db = true; });
    while (!(da && db)) {
      if (io.stopped()) io.restart();
      if (io.poll() == 0) std::this_thread::sleep_for(1ms);
    }
    REQUIRE_FALSE(ea);
    REQUIRE_FALSE(eb);
  }
};

template <typename Pred>
void pump_until(asio::io_context& io, Pred pred,
                std::chrono::milliseconds budget = 3s) {
  const auto deadline = std::chrono::steady_clock::now() + budget;
  while (!pred()) {
    if (std::chrono::steady_clock::now() > deadline) FAIL("pump_until budget exceeded");
    if (io.stopped()) io.restart();
    if (io.poll() == 0) std::this_thread::sleep_for(1ms);
  }
}

void send_bytes(SocketPair& sp, std::vector<uint8_t> bytes) {
  auto buf = std::make_shared<std::vector<uint8_t>>(std::move(bytes));
  bool done = false;
  asio::async_write(sp.b, asio::buffer(*buf),
                    [buf, &done](std::error_code ec, std::size_t) {
                      REQUIRE_FALSE(ec);
                      done = true;
                    });
  pump_until(sp.io, [&] { return done; });
}

std::optional<hsms::Frame> try_recv_frame(SocketPair& sp,
                                           std::chrono::milliseconds budget = 2s) {
  auto lenbuf = std::make_shared<std::array<uint8_t, 4>>();
  bool len_done = false;
  std::error_code rec_ec;
  asio::async_read(sp.b, asio::buffer(*lenbuf),
                   [lenbuf, &len_done, &rec_ec](std::error_code ec, std::size_t) {
                     rec_ec = ec;
                     len_done = true;
                   });
  const auto deadline = std::chrono::steady_clock::now() + budget;
  while (!len_done) {
    if (std::chrono::steady_clock::now() > deadline) return std::nullopt;
    if (sp.io.stopped()) sp.io.restart();
    if (sp.io.poll() == 0) std::this_thread::sleep_for(1ms);
  }
  if (rec_ec) return std::nullopt;
  const uint32_t len = (uint32_t((*lenbuf)[0]) << 24) | (uint32_t((*lenbuf)[1]) << 16) |
                       (uint32_t((*lenbuf)[2]) << 8) | uint32_t((*lenbuf)[3]);
  auto payload = std::make_shared<std::vector<uint8_t>>(len);
  bool payload_done = false;
  asio::async_read(sp.b, asio::buffer(*payload),
                   [payload, &payload_done](std::error_code, std::size_t) {
                     payload_done = true;
                   });
  pump_until(sp.io, [&] { return payload_done; });
  return hsms::Frame::decode(payload->data(), payload->size());
}

hsms::Timers permissive_timers() {
  hsms::Timers t;
  t.t3 = 5s; t.t6 = 5s; t.t7 = 5s; t.t8 = 5s; t.linktest = 0ms;
  return t;
}

}  // namespace

TEST_CASE("HSMS-GS passive: two sessions selected independently") {
  SocketPair sp;
  auto conn = std::make_shared<hsms::Connection>(
      std::move(sp.a), hsms::Connection::Mode::Passive,
      /*primary_device_id=*/1, permissive_timers());
  conn->add_session(2);

  std::atomic<int> sel_count_1{0};
  std::atomic<int> sel_count_2{0};
  conn->set_session_selected_handler(1, [&] { ++sel_count_1; });
  conn->set_session_selected_handler(2, [&] { ++sel_count_2; });
  conn->start();

  // Select session 1 via Select.req(session_id=1).
  send_bytes(sp, hsms::Frame(hsms::Header::control(
                       hsms::SType::SelectReq, 100, /*session_id=*/1))
                     .encode());
  auto rsp1 = try_recv_frame(sp);
  REQUIRE(rsp1.has_value());
  CHECK(rsp1->header.stype == hsms::SType::SelectRsp);
  CHECK(rsp1->header.byte3 == static_cast<uint8_t>(hsms::SelectStatus::Ok));
  pump_until(sp.io, [&] { return sel_count_1.load() == 1; });
  CHECK(conn->is_session_selected(1));
  CHECK_FALSE(conn->is_session_selected(2));

  // Select session 2.
  send_bytes(sp, hsms::Frame(hsms::Header::control(
                       hsms::SType::SelectReq, 101, /*session_id=*/2))
                     .encode());
  auto rsp2 = try_recv_frame(sp);
  REQUIRE(rsp2.has_value());
  CHECK(rsp2->header.byte3 == static_cast<uint8_t>(hsms::SelectStatus::Ok));
  pump_until(sp.io, [&] { return sel_count_2.load() == 1; });
  CHECK(conn->is_session_selected(2));

  conn->close("test done");
}

TEST_CASE("HSMS-GS: Select.req for unknown session is Rejected") {
  SocketPair sp;
  auto conn = std::make_shared<hsms::Connection>(
      std::move(sp.a), hsms::Connection::Mode::Passive,
      /*primary_device_id=*/1, permissive_timers());
  conn->add_session(2);
  conn->start();

  // Request select for session 7 — never registered.
  send_bytes(sp, hsms::Frame(hsms::Header::control(
                       hsms::SType::SelectReq, 5, /*session_id=*/7))
                     .encode());
  auto rej = try_recv_frame(sp);
  REQUIRE(rej.has_value());
  CHECK(rej->header.stype == hsms::SType::RejectReq);
  CHECK(rej->header.byte3 ==
        static_cast<uint8_t>(hsms::RejectReason::EntityNotSelected));

  conn->close("test done");
}

TEST_CASE("HSMS-GS: data routed by session_id; wrong session is rejected") {
  SocketPair sp;
  auto conn = std::make_shared<hsms::Connection>(
      std::move(sp.a), hsms::Connection::Mode::Passive,
      /*primary_device_id=*/1, permissive_timers());
  conn->add_session(2);

  std::vector<std::pair<uint16_t, std::pair<uint8_t, uint8_t>>> seen;
  conn->set_session_message_handler(1, [&](const secs2::Message& m)
                                       -> std::optional<secs2::Message> {
    seen.emplace_back(1, std::make_pair(m.stream, m.function));
    return std::nullopt;
  });
  conn->set_session_message_handler(2, [&](const secs2::Message& m)
                                       -> std::optional<secs2::Message> {
    seen.emplace_back(2, std::make_pair(m.stream, m.function));
    return std::nullopt;
  });
  conn->start();

  // Select session 1 only.
  send_bytes(sp, hsms::Frame(hsms::Header::control(
                       hsms::SType::SelectReq, 10, /*session_id=*/1))
                     .encode());
  auto rsp = try_recv_frame(sp);
  REQUIRE(rsp.has_value());
  pump_until(sp.io, [&] { return conn->is_session_selected(1); });

  // Data for session 1 — should route to session 1's handler.
  send_bytes(sp, hsms::Frame(hsms::Header::data_message(
                       /*session_id=*/1, /*stream=*/1, /*function=*/3,
                       /*reply_expected=*/false, /*sys=*/200))
                     .encode());
  pump_until(sp.io, [&] {
    for (auto& s : seen) if (s.first == 1) return true;
    return false;
  });

  // Data for session 2 (not selected) — should be Reject(EntityNotSelected).
  send_bytes(sp, hsms::Frame(hsms::Header::data_message(
                       /*session_id=*/2, /*stream=*/1, /*function=*/3,
                       /*reply_expected=*/false, /*sys=*/201))
                     .encode());
  auto rej = try_recv_frame(sp);
  REQUIRE(rej.has_value());
  CHECK(rej->header.stype == hsms::SType::RejectReq);
  CHECK(rej->header.byte3 ==
        static_cast<uint8_t>(hsms::RejectReason::EntityNotSelected));

  // Session 2's handler must NOT have fired.
  for (auto& s : seen) CHECK(s.first != 2);

  conn->close("test done");
}

TEST_CASE("HSMS-GS active: two registered sessions both end up selected") {
  // Active mode: the connection sends Select.req for each session
  // serially, waiting for each Ok before moving to the next.
  SocketPair sp;
  // Build active "host" with two sessions.
  auto host = std::make_shared<hsms::Connection>(
      std::move(sp.a), hsms::Connection::Mode::Active,
      /*primary_device_id=*/1, permissive_timers());
  host->add_session(2);
  std::atomic<int> host_sel_1{0};
  std::atomic<int> host_sel_2{0};
  host->set_session_selected_handler(1, [&] { ++host_sel_1; });
  host->set_session_selected_handler(2, [&] { ++host_sel_2; });

  // Use the *other* socket as a raw peer.  The host will issue
  // Select.req in GS mode (session_id=device_id of each session);
  // we ack each one Ok.
  host->start();

  // First Select.req should arrive for session 1.
  auto sel1 = try_recv_frame(sp);
  REQUIRE(sel1.has_value());
  CHECK(sel1->header.stype == hsms::SType::SelectReq);
  CHECK(sel1->header.session_id == 1);
  // Ack Ok.
  send_bytes(sp, hsms::Frame(hsms::Header::control(
                       hsms::SType::SelectRsp, sel1->header.system_bytes,
                       /*session_id=*/1, /*byte2=*/0,
                       /*byte3=*/static_cast<uint8_t>(hsms::SelectStatus::Ok)))
                     .encode());
  pump_until(sp.io, [&] { return host_sel_1.load() == 1; });

  // Second Select.req should arrive for session 2.
  auto sel2 = try_recv_frame(sp);
  REQUIRE(sel2.has_value());
  CHECK(sel2->header.stype == hsms::SType::SelectReq);
  CHECK(sel2->header.session_id == 2);
  send_bytes(sp, hsms::Frame(hsms::Header::control(
                       hsms::SType::SelectRsp, sel2->header.system_bytes,
                       /*session_id=*/2, /*byte2=*/0,
                       /*byte3=*/static_cast<uint8_t>(hsms::SelectStatus::Ok)))
                     .encode());
  pump_until(sp.io, [&] { return host_sel_2.load() == 1; });

  CHECK(host->is_session_selected(1));
  CHECK(host->is_session_selected(2));

  host->close("test done");
}

TEST_CASE("HSMS-SS: existing single-session API still works (backwards compat)") {
  // Sanity that the SS-only path hasn't regressed.
  SocketPair sp;
  auto conn = std::make_shared<hsms::Connection>(
      std::move(sp.a), hsms::Connection::Mode::Passive, 0, permissive_timers());
  bool selected = false;
  conn->set_selected_handler([&] { selected = true; });
  conn->start();

  // SS-style Select.req (session_id=0xFFFF).
  send_bytes(sp, hsms::Frame(hsms::Header::control(
                       hsms::SType::SelectReq, /*sys=*/1))
                     .encode());
  auto rsp = try_recv_frame(sp);
  REQUIRE(rsp.has_value());
  CHECK(rsp->header.stype == hsms::SType::SelectRsp);
  CHECK(rsp->header.byte3 == static_cast<uint8_t>(hsms::SelectStatus::Ok));
  pump_until(sp.io, [&] { return selected; });
  CHECK(conn->selected());

  conn->close("test done");
}