package tcp

import (
	"errors"
	"io"
	"net"
	"strconv"

	"github.com/panjf2000/gnet/v2"

	"spark-server/internal/transport/tcp/protocol"
)

func (s *Server) consumeProxy(st *connState) gnet.Action {
	if !s.proxyProtocol || st.sess.ProxyDone() {
		return gnet.None
	}
	avail := st.gc.InboundBuffered()
	if avail == 0 {
		return gnet.None
	}
	buf, err := st.gc.Peek(avail)
	if err != nil {
		return gnet.Close
	}
	h, n, err := protocol.ParsePrefix(buf)
	if protocol.NeedMore(err) {
		return gnet.None
	}
	if err != nil {
		return proxyFallback(st, buf)
	}
	_, _ = st.gc.Discard(n)
	st.sess.SetRealIP(headerRealIP(h, st.sess.RemoteAddr()))
	return gnet.None
}

func proxyFallback(st *connState, buf []byte) gnet.Action {
	if protocol.LooksLikeFrame(buf) {
		st.sess.MarkProxySkipped()
		return gnet.None
	}
	if protocol.MaybeFramePrefix(buf) && len(buf) < 2 {
		return gnet.None
	}
	return gnet.Close
}

func headerRealIP(h protocol.Header, fallback string) string {
	if h.SrcIP == nil {
		return fallback
	}
	return net.JoinHostPort(h.SrcIP.String(), strconv.Itoa(h.SrcPort))
}

func (s *Server) consumeFrame(st *connState) (gnet.Action, bool) {
	if st.gc.InboundBuffered() < protocol.HeaderSize {
		return gnet.None, false
	}
	header, err := st.gc.Peek(protocol.HeaderSize)
	if err != nil {
		if errors.Is(err, io.ErrShortBuffer) {
			return gnet.None, false
		}
		return gnet.Close, false
	}
	bodyLen, err := protocol.PeekBodyLength(header)
	if err != nil || bodyLen > s.maxFrame {
		return gnet.Close, false
	}
	total := protocol.HeaderSize + bodyLen
	if st.gc.InboundBuffered() < total {
		return gnet.None, false
	}
	raw, err := st.gc.Peek(total)
	if err != nil {
		if errors.Is(err, io.ErrShortBuffer) {
			return gnet.None, false
		}
		return gnet.Close, false
	}
	frame, err := protocol.Decode(raw, s.maxFrame)
	if err != nil {
		return gnet.Close, false
	}
	_, _ = st.gc.Discard(total)
	body, err := protocol.DecodeBody(frame.Body, frame.Flags, s.maxFrame)
	if err != nil {
		return gnet.Close, false
	}
	frame.Body = body
	if frame.Version != 0 && frame.Version != protocol.CurrentVersion {
		return gnet.Close, false
	}
	s.handleInbound(st, frame)
	return gnet.None, true
}