#pragma once

#include <asio.hpp>
#include <array>
#include <cstdint>
#include <deque>
#include <functional>
#include <map>
#include <memory>
#include <optional>
#include <string>
#include <system_error>

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

namespace secsgem::hsms {

// One HSMS session-set over a connected TCP socket. Drives the connection
// state machine (NOT_SELECTED -> SELECTED, per session), runs the
// SELECT/LINKTEST/SEPARATE control handshakes and the T3/T6/T7/T8 timers,
// and correlates data replies by system bytes. Single-threaded: all work
// runs on the socket's executor.
//
// Supports HSMS-SS (single-session, default — created by the constructor)
// and HSMS-GS (multi-session, E37 §11) via add_session().  In GS mode the
// Select.req carries the session's device_id in the header session_id
// field; in SS mode it carries the kControlSessionId = 0xFFFF sentinel
// per the legacy convention.
class Connection : public std::enable_shared_from_this<Connection> {
 public:
  enum class Mode { Active, Passive };
  enum class State { NotSelected, Selected };

  // Handler for an inbound primary (request) message. Return a reply Message to
  // send it back (system bytes are filled in automatically), or nullopt for none.
  using MessageHandler = std::function<std::optional<secs2::Message>(const secs2::Message&)>;
  using SelectedHandler = std::function<void()>;
  using ClosedHandler = std::function<void(const std::string& reason)>;
  using ReplyHandler = std::function<void(std::error_code, const secs2::Message&)>;
  using LogHandler = std::function<void(const std::string&)>;

  Connection(asio::ip::tcp::socket socket, Mode mode, uint16_t device_id, Timers timers);

  // SS-style setters: operate on the primary (constructor-supplied) session.
  void set_message_handler(MessageHandler h);
  void set_selected_handler(SelectedHandler h);
  void set_closed_handler(ClosedHandler h) { on_closed_ = std::move(h); }
  void set_log_handler(LogHandler h) { on_log_ = std::move(h); }

  // ---- HSMS-GS (multi-session) ------------------------------------------
  // Register an additional session.  In Active mode each registered
  // session triggers a Select.req on its device_id once the prior
  // session reaches SELECTED.  In Passive mode the session can be
  // selected by a peer's Select.req(session_id=device_id).
  void add_session(uint16_t device_id);
  void set_session_message_handler(uint16_t device_id, MessageHandler);
  void set_session_selected_handler(uint16_t device_id, SelectedHandler);
  State session_state(uint16_t device_id) const;
  bool is_session_selected(uint16_t device_id) const;

  // Multi-session send overloads — frames carry session_id = device_id.
  void send_request(uint16_t device_id, secs2::Message msg, ReplyHandler cb);
  void send_data(uint16_t device_id, secs2::Message msg);

  // Begin the read loop. Active mode also initiates the SELECT handshake;
  // Passive mode arms the T7 not-selected timer and waits for Select.req.
  void start();

  // Send a primary data message (W-bit set) and invoke `cb` with the reply or a
  // timeout/error. Must be SELECTED.
  void send_request(secs2::Message msg, ReplyHandler cb);

  // Send a data message with no reply expected.
  void send_data(secs2::Message msg);

  // Graceful teardown: send Separate.req and close.
  void separate();

  // Hard close.
  void close(const std::string& reason);

  // SS-style state getters: report the primary session's state.
  State state() const;
  bool selected() const;

  // The HSMS header of the primary currently being dispatched by the
  // message_handler.  Only non-null inside the handler.  Used so the
  // handler can capture MHEAD when it needs to emit S9F3 / S9F5 for an
  // unrecognized stream/function.
  const Header* current_header() const { return current_header_; }

  // Equipment-initiated S9F<function> primary carrying a 10-byte MHEAD.
  // Public so a wrapping message_handler can emit S9F3 / S9F5 when it
  // detects an unhandled message; Connection itself uses this for S9F7
  // (illegal data), S9F9 (T3 timeout), and S9F11 (data too long).
  void emit_s9(uint8_t function, const std::array<uint8_t, 10>& mhead);

 private:
  // --- read path ---
  void read_length();
  void on_length(std::error_code ec, std::size_t n);
  void on_payload(std::error_code ec, std::size_t n);
  void handle_frame(Frame frame);
  void handle_data(const Frame& frame);
  void handle_control(const Frame& frame);

  // --- write path ---
  void send_frame(Frame frame);
  void write_next();

  // --- handshakes & timers ---
  void send_select_req();                      // primary session
  void send_select_req(uint16_t device_id);    // explicit session (GS)
  void send_linktest_req();
  void enter_selected();                       // primary session
  void enter_selected(uint16_t device_id);     // explicit session (GS)
  void arm_t7();
  void arm_linktest();
  void start_control_transaction(SType expected_response, uint32_t system_bytes,
                                 const char* what);
  void clear_control_transaction();

  uint32_t next_system_bytes();
  void log(const std::string& msg);

  asio::ip::tcp::socket socket_;
  asio::steady_timer t6_timer_;        // control transaction timeout
  asio::steady_timer t7_timer_;        // not-selected timeout (passive)
  asio::steady_timer t8_timer_;        // intercharacter timeout (during a read)
  asio::steady_timer linktest_timer_;  // periodic linktest interval

  Mode mode_;
  uint16_t device_id_;                  // primary session's device_id
  Timers timers_;
  bool closed_ = false;

  // Per-session state.  In SS mode this map has a single entry keyed
  // by device_id_.  In GS mode add_session() inserts more entries.
  struct SessionSlot {
    uint16_t device_id;
    State state = State::NotSelected;
    MessageHandler on_message;
    SelectedHandler on_selected;
  };
  std::map<uint16_t, SessionSlot> sessions_;

  // Active mode walk-list: device_ids still awaiting their initial
  // Select.req when running multi-session.  Empty for SS.
  std::deque<uint16_t> pending_active_selects_;

  bool gs_mode() const { return sessions_.size() > 1; }
  SessionSlot* find_session(uint16_t device_id) {
    auto it = sessions_.find(device_id);
    return it == sessions_.end() ? nullptr : &it->second;
  }
  const SessionSlot* find_session(uint16_t device_id) const {
    auto it = sessions_.find(device_id);
    return it == sessions_.end() ? nullptr : &it->second;
  }
  // Pick which session an inbound Select.req with session_id=`sid` is
  // targeting.  SS legacy: any session_id (often 0xFFFF) maps to the
  // single primary session.  GS: must match a registered device_id.
  SessionSlot* resolve_select_target(uint16_t sid);
  // Pick which session an inbound data frame with session_id=`sid`
  // is for.  SS: primary.  GS: must match a registered device_id.
  SessionSlot* resolve_data_target(uint16_t sid);

  std::array<uint8_t, kLengthPrefixSize> len_buf_{};
  std::vector<uint8_t> payload_;

  std::deque<std::vector<uint8_t>> write_queue_;
  bool writing_ = false;
  bool close_after_flush_ = false;
  std::string close_reason_;

  uint32_t next_system_bytes_ = 1;

  // outstanding data request transactions, keyed by system bytes.
  // We track the EXPECTED reply stream/function so we can match exactly
  // instead of inferring "this is a reply" from function parity.
  struct PendingRequest {
    uint8_t expected_stream;
    uint8_t expected_function;
    ReplyHandler cb;
    std::shared_ptr<asio::steady_timer> t3;
  };
  std::map<uint32_t, PendingRequest> pending_requests_;

  // single outstanding control transaction (select / linktest)
  struct PendingControl {
    SType expected_response;
    uint32_t system_bytes;
  };
  std::optional<PendingControl> pending_control_;

  const Header* current_header_ = nullptr;  // set only inside dispatch

  ClosedHandler on_closed_;
  LogHandler on_log_;
};

// Error category for HSMS protocol-level failures surfaced through ReplyHandler.
enum class Error {
  Timeout = 1,
  Closed = 2,
  IllegalData = 3,  // peer's reply body couldn't be decoded; we also send S9F7
};
std::error_code make_error(Error e);

}  // namespace secsgem::hsms