#include "secsgem/hsms/connection.hpp"

#include "secsgem/secs2/codec.hpp"

#include <exception>

namespace secsgem::hsms {
namespace {

class HsmsCategory : public std::error_category {
 public:
  const char* name() const noexcept override { return "hsms"; }
  std::string message(int ev) const override {
    switch (static_cast<Error>(ev)) {
      case Error::Timeout:     return "HSMS transaction timeout";
      case Error::Closed:      return "HSMS connection closed";
      case Error::IllegalData: return "peer sent a body that failed to decode";
    }
    return "unknown HSMS error";
  }
};

const HsmsCategory g_category;

constexpr std::size_t kMaxFrameLength = 16 * 1024 * 1024;

}  // namespace

std::error_code make_error(Error e) {
  return {static_cast<int>(e), g_category};
}

Connection::Connection(asio::ip::tcp::socket socket, Mode mode, uint16_t device_id,
                       Timers timers)
    : socket_(std::move(socket)),
      t6_timer_(socket_.get_executor()),
      t7_timer_(socket_.get_executor()),
      t8_timer_(socket_.get_executor()),
      linktest_timer_(socket_.get_executor()),
      mode_(mode),
      device_id_(device_id),
      timers_(timers) {
  SessionSlot primary;
  primary.device_id = device_id;
  sessions_.emplace(device_id, std::move(primary));
}

void Connection::set_message_handler(MessageHandler h) {
  if (auto* s = find_session(device_id_)) s->on_message = std::move(h);
}
void Connection::set_selected_handler(SelectedHandler h) {
  if (auto* s = find_session(device_id_)) s->on_selected = std::move(h);
}

Connection::State Connection::state() const {
  auto* s = find_session(device_id_);
  return s ? s->state : State::NotSelected;
}
bool Connection::selected() const { return state() == State::Selected; }

void Connection::add_session(uint16_t device_id) {
  if (sessions_.count(device_id)) return;
  SessionSlot s;
  s.device_id = device_id;
  sessions_.emplace(device_id, std::move(s));
  if (mode_ == Mode::Active) {
    pending_active_selects_.push_back(device_id);
  }
}

void Connection::set_session_message_handler(uint16_t device_id, MessageHandler h) {
  if (auto* s = find_session(device_id)) s->on_message = std::move(h);
}
void Connection::set_session_selected_handler(uint16_t device_id, SelectedHandler h) {
  if (auto* s = find_session(device_id)) s->on_selected = std::move(h);
}
Connection::State Connection::session_state(uint16_t device_id) const {
  auto* s = find_session(device_id);
  return s ? s->state : State::NotSelected;
}
bool Connection::is_session_selected(uint16_t device_id) const {
  return session_state(device_id) == State::Selected;
}

Connection::SessionSlot* Connection::resolve_select_target(uint16_t sid) {
  if (gs_mode()) return find_session(sid);
  // SS legacy: session_id is typically 0xFFFF in Select.req but some
  // peers also send the device_id.  Either way, route to the primary.
  return find_session(device_id_);
}

Connection::SessionSlot* Connection::resolve_data_target(uint16_t sid) {
  if (gs_mode()) return find_session(sid);
  return find_session(device_id_);
}

void Connection::start() {
  read_length();
  if (mode_ == Mode::Active) {
    if (gs_mode()) send_select_req(device_id_);
    else           send_select_req();
  } else {
    arm_t7();
  }
}

// --- read path ------------------------------------------------------------

void Connection::read_length() {
  if (closed_) return;
  auto self = shared_from_this();
  asio::async_read(socket_, asio::buffer(len_buf_),
                   [this, self](std::error_code ec, std::size_t n) { on_length(ec, n); });
}

void Connection::on_length(std::error_code ec, std::size_t) {
  if (closed_) return;
  if (ec) {
    close("read length: " + ec.message());
    return;
  }
  const uint32_t len = (static_cast<uint32_t>(len_buf_[0]) << 24) |
                       (static_cast<uint32_t>(len_buf_[1]) << 16) |
                       (static_cast<uint32_t>(len_buf_[2]) << 8) |
                       static_cast<uint32_t>(len_buf_[3]);
  if (len < kHeaderSize) {
    close("invalid frame length " + std::to_string(len));
    return;
  }
  if (len > kMaxFrameLength) {
    // We can't read the actual offending header (the message body is too
    // big to safely buffer), so synthesize a 10-byte MHEAD whose first 4
    // bytes are the offending length prefix and the rest are zero.  The
    // host sees S9F11, the connection then drains the write queue and
    // closes.
    std::array<uint8_t, 10> mhead{};
    mhead[0] = len_buf_[0]; mhead[1] = len_buf_[1];
    mhead[2] = len_buf_[2]; mhead[3] = len_buf_[3];
    log("frame too large: " + std::to_string(len) + "; emitting S9F11 and closing");
    emit_s9(11, mhead);
    close_after_flush_ = true;
    close_reason_ = "frame too large";
    return;
  }

  payload_.resize(len);

  t8_timer_.expires_after(timers_.t8);
  auto self = shared_from_this();
  t8_timer_.async_wait([this, self](std::error_code tec) {
    if (!tec && !closed_) close("T8 intercharacter timeout");
  });
  asio::async_read(socket_, asio::buffer(payload_),
                   [this, self](std::error_code rec, std::size_t n) { on_payload(rec, n); });
}

void Connection::on_payload(std::error_code ec, std::size_t) {
  if (closed_) return;
  t8_timer_.cancel();
  if (ec) {
    close("read payload: " + ec.message());
    return;
  }
  Frame frame;
  try {
    frame = Frame::decode(payload_.data(), payload_.size());
  } catch (const std::exception& e) {
    close(std::string("decode: ") + e.what());
    return;
  }
  // E37 §7.7: a non-SECS-II PType must be rejected with Reject.req.
  // Convention (matches the EntityNotSelected emission below): the
  // offending header byte is echoed in byte2 and the reason code goes
  // in byte3.
  if (frame.header.ptype != kPTypeSecsII) {
    log("!! unsupported PType=" + std::to_string(frame.header.ptype) +
        "; emitting Reject.req");
    send_frame(Frame(Header::control(
        SType::RejectReq, frame.header.system_bytes, frame.header.session_id,
        frame.header.ptype,
        static_cast<uint8_t>(RejectReason::PtypeNotSupported))));
    read_length();
    return;
  }
  handle_frame(std::move(frame));
  read_length();
}

void Connection::handle_frame(Frame frame) {
  if (frame.header.stype == SType::Data) {
    handle_data(frame);
    return;
  }
  // E37 §7.7: any SType the receiver does not implement must be answered
  // with Reject.req (reason=StypeNotSupported). Defined SType values are
  // {1,2,3,4,5,6,7,9}; anything else here is unsupported.
  switch (frame.header.stype) {
    case SType::SelectReq:
    case SType::SelectRsp:
    case SType::DeselectReq:
    case SType::DeselectRsp:
    case SType::LinktestReq:
    case SType::LinktestRsp:
    case SType::RejectReq:
    case SType::SeparateReq:
      handle_control(frame);
      return;
    case SType::Data:
      return;  // already handled above
  }
  const uint8_t bad_stype = static_cast<uint8_t>(frame.header.stype);
  log("!! unsupported SType=" + std::to_string(bad_stype) +
      "; emitting Reject.req");
  send_frame(Frame(Header::control(
      SType::RejectReq, frame.header.system_bytes, frame.header.session_id,
      bad_stype,
      static_cast<uint8_t>(RejectReason::StypeNotSupported))));
}

void Connection::handle_data(const Frame& frame) {
  const Header& h = frame.header;

  // Reply correlation: match on (system_bytes, stream, function) exactly.
  // System bytes are connection-scope (E37 §8.3); replies route by system
  // bytes regardless of which session they came in on.
  auto it = pending_requests_.find(h.system_bytes);
  if (it != pending_requests_.end() &&
      it->second.expected_stream == h.stream() &&
      it->second.expected_function == h.function()) {
    secs2::Message reply;
    try {
      reply = secs2::Message::from_body(h.stream(), h.function(), h.w_bit(), frame.body);
    } catch (const std::exception& e) {
      // Malformed reply: notify the peer with S9F7 and surface IllegalData
      // to the caller right now (rather than making them wait out T3).
      log(std::string("reply body decode: ") + e.what() + "; emitting S9F7");
      emit_s9(7, h.encode());
      auto cb = std::move(it->second.cb);
      it->second.t3->cancel();
      pending_requests_.erase(it);
      cb(make_error(Error::IllegalData), {});
      return;
    }
    auto cb = std::move(it->second.cb);
    it->second.t3->cancel();
    pending_requests_.erase(it);
    log("<- " + h.describe());
    cb({}, reply);
    return;
  }

  // System bytes match but stream/function disagree: protocol violation
  // (peer replied with the wrong SxFy). Surface this as a real diagnostic
  // rather than silently falling through to the primary handler.
  if (it != pending_requests_.end()) {
    log("!! unexpected " + h.describe() + " for pending sys=" +
        std::to_string(h.system_bytes) + " (expected S" +
        std::to_string(it->second.expected_stream) + "F" +
        std::to_string(it->second.expected_function) + ")");
    // Treat the data as a primary message anyway so the application
    // handler can deal with it; the pending request stays open until T3.
  }

  // Primary message; route by header.session_id (GS) or to primary (SS).
  // Reject if no matching session OR if that session isn't SELECTED.
  SessionSlot* target = resolve_data_target(h.session_id);
  if (!target || target->state != State::Selected) {
    send_frame(Frame(Header::control(SType::RejectReq, h.system_bytes, kControlSessionId,
                                     static_cast<uint8_t>(SType::Data),
                                     static_cast<uint8_t>(RejectReason::EntityNotSelected))));
    return;
  }

  secs2::Message msg;
  try {
    msg = secs2::Message::from_body(h.stream(), h.function(), h.w_bit(), frame.body);
  } catch (const std::exception& e) {
    // E5: a malformed body in a primary is reported via S9F7 (Illegal Data).
    // The connection stays up — the peer may send corrected data next.
    log(std::string("message body decode: ") + e.what() + "; emitting S9F7");
    emit_s9(7, h.encode());
    return;
  }

  log("<- " + h.describe());
  if (!target->on_message) return;

  // Expose the originating header to the handler in case it needs to emit
  // an S9F3 / S9F5 in response.  Cleared on the way out.
  current_header_ = &h;
  auto reply = target->on_message(msg);
  current_header_ = nullptr;

  if (reply) {
    Frame out(Header::data_message(target->device_id, reply->stream, reply->function,
                                   reply->reply_expected, h.system_bytes),
              reply->encode_body());
    log("-> " + out.header.describe());
    send_frame(std::move(out));
  }
}

void Connection::handle_control(const Frame& frame) {
  const Header& h = frame.header;
  switch (h.stype) {
    case SType::SelectReq: {
      // Route by header.session_id (GS) or to primary (SS legacy).
      SessionSlot* tgt = resolve_select_target(h.session_id);
      if (!tgt) {
        // GS: no such session on this connection.
        log("<- Select.req for unknown session " + std::to_string(h.session_id));
        send_frame(Frame(Header::control(
            SType::RejectReq, h.system_bytes, h.session_id,
            static_cast<uint8_t>(SType::SelectReq),
            static_cast<uint8_t>(RejectReason::EntityNotSelected))));
        break;
      }
      const auto status = (tgt->state == State::Selected)
                              ? SelectStatus::AlreadyActive
                              : SelectStatus::Ok;
      log(std::string("<- Select.req") +
          (status == SelectStatus::AlreadyActive ? " (already SELECTED)" : ""));
      send_frame(Frame(Header::control(SType::SelectRsp, h.system_bytes, h.session_id, 0,
                                       static_cast<uint8_t>(status))));
      log(std::string("-> Select.rsp (") +
          (status == SelectStatus::Ok ? "Ok" : "AlreadyActive") + ")");
      if (status == SelectStatus::Ok) enter_selected(tgt->device_id);
      break;
    }
    case SType::SelectRsp: {
      if (pending_control_ && pending_control_->expected_response == SType::SelectRsp &&
          pending_control_->system_bytes == h.system_bytes) {
        clear_control_transaction();
        if (h.byte3 == static_cast<uint8_t>(SelectStatus::Ok)) {
          log("<- Select.rsp (Ok)");
          // Identify which session this rsp completes by session_id; SS
          // legacy mode uses 0xFFFF on the wire, default to primary.
          uint16_t selected = gs_mode() ? h.session_id : device_id_;
          enter_selected(selected);
          // Active mode walk-list: kick off next pending select.
          if (mode_ == Mode::Active && !pending_active_selects_.empty()) {
            uint16_t next = pending_active_selects_.front();
            pending_active_selects_.pop_front();
            send_select_req(next);
          }
        } else {
          close("Select rejected, status=" + std::to_string(h.byte3));
        }
      }
      break;
    }
    case SType::DeselectReq: {
      SessionSlot* tgt = resolve_select_target(h.session_id);
      const auto status = (tgt && tgt->state == State::Selected)
                              ? DeselectStatus::Ok
                              : DeselectStatus::NotEstablished;
      log(std::string("<- Deselect.req") +
          (status == DeselectStatus::NotEstablished ? " (not SELECTED)" : ""));
      send_frame(Frame(Header::control(SType::DeselectRsp, h.system_bytes, h.session_id, 0,
                                       static_cast<uint8_t>(status))));
      if (status == DeselectStatus::Ok && tgt) {
        tgt->state = State::NotSelected;
        // Re-arm T7 only if NO sessions are selected.
        bool any_selected = false;
        for (auto& kv : sessions_) if (kv.second.state == State::Selected) { any_selected = true; break; }
        if (!any_selected) arm_t7();
      }
      break;
    }
    case SType::DeselectRsp: {
      if (pending_control_ && pending_control_->expected_response == SType::DeselectRsp &&
          pending_control_->system_bytes == h.system_bytes) {
        clear_control_transaction();
        uint16_t target = gs_mode() ? h.session_id : device_id_;
        if (auto* s = find_session(target)) s->state = State::NotSelected;
      }
      break;
    }
    case SType::LinktestReq: {
      send_frame(Frame(Header::control(SType::LinktestRsp, h.system_bytes)));
      break;
    }
    case SType::LinktestRsp: {
      if (pending_control_ && pending_control_->expected_response == SType::LinktestRsp &&
          pending_control_->system_bytes == h.system_bytes) {
        clear_control_transaction();
        arm_linktest();
      }
      break;
    }
    case SType::SeparateReq: {
      log("<- Separate.req");
      close("peer separated");
      break;
    }
    case SType::RejectReq: {
      log("<- Reject.req reason=" + std::to_string(h.byte3));
      break;
    }
    case SType::Data:
      break;  // unreachable: SType::Data is dispatched to handle_data
  }
}

// --- write path -----------------------------------------------------------

void Connection::send_frame(Frame frame) {
  if (closed_) return;
  write_queue_.push_back(frame.encode());
  if (!writing_) write_next();
}

void Connection::write_next() {
  if (write_queue_.empty()) {
    writing_ = false;
    if (close_after_flush_) close(close_reason_);
    return;
  }
  writing_ = true;
  auto self = shared_from_this();
  asio::async_write(socket_, asio::buffer(write_queue_.front()),
                    [this, self](std::error_code ec, std::size_t) {
                      write_queue_.pop_front();
                      if (ec) {
                        close("write: " + ec.message());
                        return;
                      }
                      write_next();
                    });
}

// --- public send API ------------------------------------------------------

void Connection::send_request(secs2::Message msg, ReplyHandler cb) {
  send_request(device_id_, std::move(msg), std::move(cb));
}

void Connection::send_request(uint16_t device_id, secs2::Message msg, ReplyHandler cb) {
  if (closed_) {
    cb(make_error(Error::Closed), {});
    return;
  }
  const uint32_t sys = next_system_bytes();
  Frame f(Header::data_message(device_id, msg.stream, msg.function, true, sys),
          msg.encode_body());

  auto t3 = std::make_shared<asio::steady_timer>(socket_.get_executor());
  t3->expires_after(timers_.t3);
  pending_requests_.emplace(sys, PendingRequest{
      /*expected_stream=*/msg.stream,
      /*expected_function=*/static_cast<uint8_t>(msg.function + 1),
      std::move(cb), t3});

  auto self = shared_from_this();
  t3->async_wait([this, self, sys, device_id](std::error_code ec) {
    if (ec) return;
    auto it = pending_requests_.find(sys);
    if (it == pending_requests_.end()) return;
    const uint8_t req_stream   = it->second.expected_stream;
    const uint8_t req_function = static_cast<uint8_t>(it->second.expected_function - 1);
    Header h = Header::data_message(device_id, req_stream, req_function, true, sys);
    emit_s9(9, h.encode());
    auto cb2 = std::move(it->second.cb);
    pending_requests_.erase(it);
    cb2(make_error(Error::Timeout), {});
  });

  log("-> " + f.header.describe());
  send_frame(std::move(f));
}

void Connection::send_data(secs2::Message msg) {
  send_data(device_id_, std::move(msg));
}

void Connection::send_data(uint16_t device_id, secs2::Message msg) {
  if (closed_) return;
  const uint32_t sys = next_system_bytes();
  Frame f(Header::data_message(device_id, msg.stream, msg.function, msg.reply_expected, sys),
          msg.encode_body());
  log("-> " + f.header.describe());
  send_frame(std::move(f));
}

void Connection::separate() {
  if (closed_) return;
  Frame f(Header::control(SType::SeparateReq, next_system_bytes(), device_id_));
  log("-> Separate.req");
  close_after_flush_ = true;
  close_reason_ = "separated";
  send_frame(std::move(f));
}

void Connection::close(const std::string& reason) {
  if (closed_) return;
  closed_ = true;

  std::error_code ignore;
  t6_timer_.cancel();
  t7_timer_.cancel();
  t8_timer_.cancel();
  linktest_timer_.cancel();
  socket_.shutdown(asio::ip::tcp::socket::shutdown_both, ignore);
  socket_.close(ignore);

  for (auto& [sys, pr] : pending_requests_) {
    pr.t3->cancel();
    pr.cb(make_error(Error::Closed), {});
  }
  pending_requests_.clear();

  log("xx CLOSED: " + reason);
  if (on_closed_) on_closed_(reason);
}

// --- handshakes & timers --------------------------------------------------

void Connection::send_select_req() {
  // SS legacy: Select.req uses 0xFFFF session id.  Used unchanged when
  // running with a single registered session.
  const uint32_t sys = next_system_bytes();
  log("-> Select.req");
  send_frame(Frame(Header::control(SType::SelectReq, sys)));
  start_control_transaction(SType::SelectRsp, sys, "Select");
}

void Connection::send_select_req(uint16_t device_id) {
  // GS: Select.req carries the session's device_id in the session_id
  // field (E37 §11.4.2).
  const uint32_t sys = next_system_bytes();
  log("-> Select.req session=" + std::to_string(device_id));
  send_frame(Frame(Header::control(SType::SelectReq, sys, device_id)));
  start_control_transaction(SType::SelectRsp, sys, "Select");
}

void Connection::send_linktest_req() {
  const uint32_t sys = next_system_bytes();
  send_frame(Frame(Header::control(SType::LinktestReq, sys)));
  start_control_transaction(SType::LinktestRsp, sys, "Linktest");
}

void Connection::enter_selected() {
  enter_selected(device_id_);
}

void Connection::enter_selected(uint16_t device_id) {
  auto* s = find_session(device_id);
  if (!s) return;
  if (s->state == State::Selected) return;
  s->state = State::Selected;
  t7_timer_.cancel();
  log("== SELECTED session=" + std::to_string(device_id) + " ==");
  arm_linktest();
  if (s->on_selected) s->on_selected();
}

void Connection::arm_t7() {
  if (mode_ != Mode::Passive) return;
  if (timers_.t7.count() <= 0) return;
  t7_timer_.expires_after(timers_.t7);
  auto self = shared_from_this();
  t7_timer_.async_wait([this, self](std::error_code ec) {
    if (!ec && !closed_) {
      // Fire only if NO session has been selected — T7 is a single
      // connection-scope timer per E37 §10, not per-session.
      bool any = false;
      for (auto& kv : sessions_) if (kv.second.state == State::Selected) { any = true; break; }
      if (!any) close("T7 not-selected timeout");
    }
  });
}

void Connection::arm_linktest() {
  if (timers_.linktest.count() <= 0) return;
  // Run linktest while any session is selected — that proves the
  // underlying TCP path is alive regardless of which session uses it.
  bool any = false;
  for (auto& kv : sessions_) if (kv.second.state == State::Selected) { any = true; break; }
  if (!any) return;
  linktest_timer_.expires_after(timers_.linktest);
  auto self = shared_from_this();
  linktest_timer_.async_wait([this, self](std::error_code ec) {
    if (ec || closed_) return;
    bool any2 = false;
    for (auto& kv : sessions_) if (kv.second.state == State::Selected) { any2 = true; break; }
    if (any2) send_linktest_req();
  });
}

void Connection::start_control_transaction(SType expected_response, uint32_t system_bytes,
                                           const char* what) {
  pending_control_ = PendingControl{expected_response, system_bytes};
  t6_timer_.expires_after(timers_.t6);
  auto self = shared_from_this();
  std::string what_str = what;
  t6_timer_.async_wait([this, self, what_str](std::error_code ec) {
    if (!ec && !closed_ && pending_control_) close("T6 timeout on " + what_str);
  });
}

void Connection::clear_control_transaction() {
  pending_control_.reset();
  t6_timer_.cancel();
}

void Connection::emit_s9(uint8_t function, const std::array<uint8_t, 10>& mhead) {
  if (closed_) return;
  const uint32_t sys = next_system_bytes();
  // S9F1/F3/F5/F7/F9/F11 all carry the offending MHEAD as <B 10> in the body.
  secs2::Item body_item =
      secs2::Item::binary(std::vector<uint8_t>(mhead.begin(), mhead.end()));
  Frame f(Header::data_message(device_id_, /*stream=*/9, function,
                               /*w=*/false, sys),
          secs2::encode(body_item));
  log("-> S9F" + std::to_string(function) + " (sys=" + std::to_string(sys) + ")");
  send_frame(std::move(f));
}

uint32_t Connection::next_system_bytes() {
  uint32_t s = next_system_bytes_++;
  if (next_system_bytes_ == 0) next_system_bytes_ = 1;
  return s;
}

void Connection::log(const std::string& msg) {
  if (on_log_) on_log_(msg);
}

}  // namespace secsgem::hsms