package tcp

import (
	"context"

	"spark-server/internal/pkg/errs"
	"spark-server/internal/pkg/reqctx"
	"spark-server/internal/transport/tcp/protocol"
	"spark-server/internal/transport/tcp/session"
)

func (s *Server) handleInbound(st *connState, frame protocol.Frame) {
	st.sess.Touch()
	if frame.Command != protocol.CmdHeartbeat && frame.Command != protocol.CmdPushAck && !st.limit.Allow() {
		s.metrics.Inc("tcp_rate_limited_total")
		s.replyErr(st, frame, errs.ErrRateLimited)
		return
	}
	if err := ensureAuth(st.sess, frame.Command); err != nil {
		s.replyErr(st, frame, err)
		return
	}
	switch frame.Command {
	case protocol.CmdHeartbeat:
		s.writeEncoded(st, replyFrame(protocol.CmdHeartbeat, frame.SequenceID, frame.RequestID, nil, 0))
	case protocol.CmdPushAck:
		return
	default:
		s.submitBusiness(st, frame)
	}
}

func ensureAuth(conn *session.Conn, cmd uint16) error {
	switch cmd {
	case protocol.CmdAuth, protocol.CmdHeartbeat, protocol.CmdPushAck:
		return nil
	default:
		if !conn.Authenticated() {
			return errs.ErrNotAuthenticated
		}
		return nil
	}
}

func (s *Server) submitBusiness(st *connState, frame protocol.Frame) {
	ctx := reqctx.With(context.Background(), st.sess.RequestContext())
	copied := frame
	err := s.pool.Submit(ctx, func() {
		s.runHandler(ctx, st.sess, copied)
	})
	if err == errs.ErrBusy {
		s.metrics.Inc("taskpool_rejected_total")
		s.replyErr(st, frame, errs.ErrBusy)
	}
}

func (s *Server) runHandler(ctx context.Context, conn *session.Conn, frame protocol.Frame) {
	ctx, cancel := context.WithTimeout(ctx, s.handlerTimeout)
	defer cancel()
	resp, err := s.router.Dispatch(ctx, conn, frame)
	if err != nil {
		_ = conn.Send(&session.Packet{
			Command: frame.Command, SequenceID: frame.SequenceID, RequestID: frame.RequestID,
			Body: protocol.EncodeError(errorCode(err)),
		})
		return
	}
	if resp == nil {
		return
	}
	_ = conn.Send(&session.Packet{
		Command: resp.Command, SequenceID: resp.SequenceID, RequestID: resp.RequestID, Flags: resp.Flags, Body: resp.Body,
	})
}

func (s *Server) replyErr(st *connState, frame protocol.Frame, err error) {
	body := protocol.EncodeError(errorCode(err))
	s.writeEncoded(st, replyFrame(frame.Command, frame.SequenceID, frame.RequestID, body, 0))
}

func (s *Server) writeEncoded(st *connState, frame protocol.Frame) {
	body, flags, err := protocol.PrepareBody(frame.Body, s.compressMin)
	if err != nil {
		return
	}
	frame.Body = body
	frame.Flags |= flags
	raw, err := protocol.Encode(frame)
	if err != nil {
		return
	}
	_ = st.gc.AsyncWrite(raw, nil)
}