#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 over a connected TCP socket. Drives the connection state
// machine (NOT_SELECTED -> SELECTED), 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.
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);

  void set_message_handler(MessageHandler h) { on_message_ = std::move(h); }
  void set_selected_handler(SelectedHandler h) { on_selected_ = std::move(h); }
  void set_closed_handler(ClosedHandler h) { on_closed_ = std::move(h); }
  void set_log_handler(LogHandler h) { on_log_ = std::move(h); }

  // 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);

  State state() const { return state_; }
  bool selected() const { return state_ == State::Selected; }

  // 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();
  void send_linktest_req();
  void enter_selected();
  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_;
  Timers timers_;
  State state_ = State::NotSelected;
  bool closed_ = false;

  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 on_message_ dispatch

  MessageHandler on_message_;
  SelectedHandler on_selected_;
  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