package protocol

import (
	"bufio"
	"encoding/binary"
	"errors"
	"io"
	"net"
	"strings"

	"spark-server/internal/pkg/errs"
)

const maxProxyLine = 108

var (
	errNeedMore  = errors.New("proxy header incomplete")
	proxyV2Sig   = []byte{0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A}
	proxyV1Magic = []byte("PROXY ")
)

type Header struct {
	SrcIP   net.IP
	SrcPort int
}

func NeedMore(err error) bool {
	return errors.Is(err, errNeedMore)
}

// ParsePrefix consumes a PROXY v1/v2 header from buf. n is bytes to discard.
func ParsePrefix(buf []byte) (h Header, n int, err error) {
	if len(buf) == 0 {
		return Header{}, 0, errNeedMore
	}
	if isPartialPrefix(buf, proxyV1Magic) || isPartialPrefix(buf, proxyV2Sig) {
		return Header{}, 0, errNeedMore
	}
	if hasPrefix(buf, proxyV1Magic) {
		return parseV1(buf)
	}
	if hasPrefix(buf, proxyV2Sig) {
		return parseV2(buf)
	}
	return Header{}, 0, errs.ErrProxyHeader
}

func isPartialPrefix(buf, magic []byte) bool {
	if len(buf) >= len(magic) {
		return false
	}
	return hasPrefix(magic, buf)
}

func parseV1(buf []byte) (Header, int, error) {
	idx := indexCRLF(buf)
	if idx < 0 {
		if len(buf) > maxProxyLine {
			return Header{}, 0, errs.ErrProxyHeader
		}
		return Header{}, 0, errNeedMore
	}
	line := string(buf[:idx])
	h, err := parseV1Line(line)
	if err != nil {
		return Header{}, 0, err
	}
	return h, idx + 2, nil
}

func parseV1Line(line string) (Header, error) {
	fields := strings.Fields(strings.TrimSpace(line))
	if len(fields) < 2 || fields[0] != "PROXY" {
		return Header{}, errs.ErrProxyHeader
	}
	if fields[1] == "UNKNOWN" {
		return Header{}, nil
	}
	if len(fields) < 6 {
		return Header{}, errs.ErrProxyHeader
	}
	ip := net.ParseIP(fields[2])
	if ip == nil {
		return Header{}, errs.ErrProxyHeader
	}
	port, ok := atoi(fields[4])
	if !ok {
		return Header{}, errs.ErrProxyHeader
	}
	return Header{SrcIP: ip, SrcPort: port}, nil
}

func parseV2(buf []byte) (Header, int, error) {
	if len(buf) < 16 {
		return Header{}, 0, errNeedMore
	}
	addrLen := int(binary.BigEndian.Uint16(buf[14:16]))
	total := 16 + addrLen
	if len(buf) < total {
		return Header{}, 0, errNeedMore
	}
	h := Header{}
	fam := buf[13] >> 4
	if fam == 1 && addrLen >= 12 { // IPv4
		h.SrcIP = net.IPv4(buf[16], buf[17], buf[18], buf[19])
		h.SrcPort = int(binary.BigEndian.Uint16(buf[24:26]))
	}
	if fam == 2 && addrLen >= 36 { // IPv6
		h.SrcIP = net.IP(append([]byte(nil), buf[16:32]...))
		h.SrcPort = int(binary.BigEndian.Uint16(buf[48:50]))
	}
	return h, total, nil
}

func ReadHeader(r *bufio.Reader) (Header, error) {
	var acc []byte
	for {
		b, err := r.ReadByte()
		if err != nil {
			return Header{}, err
		}
		acc = append(acc, b)
		if len(acc) > maxProxyLine+16 {
			return Header{}, errs.ErrProxyHeader
		}
		h, n, err := ParsePrefix(acc)
		if NeedMore(err) {
			continue
		}
		if err != nil {
			return Header{}, err
		}
		if n < len(acc) {
			for i := len(acc) - 1; i >= n; i-- {
				_ = r.UnreadByte()
			}
		}
		return h, nil
	}
}

func Discard(r io.Reader) error {
	_, err := io.Copy(io.Discard, r)
	return err
}

func hasPrefix(b, pre []byte) bool {
	if len(b) < len(pre) {
		return false
	}
	for i := range pre {
		if b[i] != pre[i] {
			return false
		}
	}
	return true
}

func indexCRLF(b []byte) int {
	for i := 0; i+1 < len(b); i++ {
		if b[i] == '\r' && b[i+1] == '\n' {
			return i
		}
	}
	return -1
}

func atoi(s string) (int, bool) {
	n := 0
	if s == "" {
		return 0, false
	}
	for _, c := range s {
		if c < '0' || c > '9' {
			return 0, false
		}
		n = n*10 + int(c-'0')
	}
	return n, true
}