package tcp

import (
	"unicode/utf8"

	"github.com/apernet/OpenGFW/analyzer"
	"github.com/apernet/OpenGFW/analyzer/utils"
)

const (
	minecraftMaxHandshakePacketLen = 2048
	minecraftMaxServerAddrLen      = 255 * 4

	minecraftPacketHandshake = 0

	minecraftNextStateStatus   = 1
	minecraftNextStateLogin    = 2
	minecraftNextStateTransfer = 3
)

var _ analyzer.TCPAnalyzer = (*MinecraftAnalyzer)(nil)

type MinecraftAnalyzer struct{}

func (a *MinecraftAnalyzer) Name() string {
	return "minecraft"
}

func (a *MinecraftAnalyzer) Limit() int {
	return minecraftMaxHandshakePacketLen + 5
}

func (a *MinecraftAnalyzer) NewTCP(info analyzer.TCPInfo, logger analyzer.Logger) analyzer.TCPStream {
	return newMinecraftStream(logger)
}

type minecraftStream struct {
	logger analyzer.Logger
	reqBuf *utils.ByteBuffer
}

func newMinecraftStream(logger analyzer.Logger) *minecraftStream {
	return &minecraftStream{logger: logger, reqBuf: &utils.ByteBuffer{}}
}

func (s *minecraftStream) Feed(rev, start, end bool, skip int, data []byte) (u *analyzer.PropUpdate, done bool) {
	if skip != 0 {
		return minecraftResult(false, false, nil), true
	}
	if len(data) == 0 {
		return nil, false
	}
	if rev {
		return nil, false
	}

	s.reqBuf.Append(data)
	m, needMore, valid := parseMinecraftHandshake(s.reqBuf.Buf)
	if needMore {
		return minecraftResult(false, true, nil), false
	}
	if !valid {
		return minecraftResult(false, false, nil), true
	}

	return minecraftResult(true, false, m), true
}

func (s *minecraftStream) Close(limited bool) *analyzer.PropUpdate {
	s.reqBuf.Reset()
	return minecraftResult(false, false, nil)
}

func minecraftResult(yes, candidate bool, m analyzer.PropMap) *analyzer.PropUpdate {
	if m == nil {
		m = analyzer.PropMap{}
	}
	m["yes"] = yes
	m["candidate"] = candidate
	return &analyzer.PropUpdate{
		Type: analyzer.PropUpdateReplace,
		M:    m,
	}
}

func parseMinecraftHandshake(buf []byte) (m analyzer.PropMap, needMore, valid bool) {
	packetLen, off, ok, invalid := readMinecraftVarInt(buf, 0)
	if invalid || packetLen <= 0 || packetLen > minecraftMaxHandshakePacketLen {
		return nil, false, false
	}
	if !ok {
		return nil, true, false
	}

	packetEnd := off + int(packetLen)
	if len(buf) < packetEnd {
		return nil, true, false
	}

	packet := buf[off:packetEnd]
	packetID, pos, ok, invalid := readMinecraftVarInt(packet, 0)
	if invalid || !ok || packetID != minecraftPacketHandshake {
		return nil, false, false
	}

	protocolVersion, pos, ok, invalid := readMinecraftVarInt(packet, pos)
	if invalid || !ok {
		return nil, false, false
	}

	addrLen, pos, ok, invalid := readMinecraftVarInt(packet, pos)
	if invalid || !ok || addrLen <= 0 || addrLen > minecraftMaxServerAddrLen {
		return nil, false, false
	}

	addrEnd := pos + int(addrLen)
	if addrEnd+2 > len(packet) {
		return nil, false, false
	}
	addrBytes := packet[pos:addrEnd]
	if !utf8.Valid(addrBytes) {
		return nil, false, false
	}
	pos = addrEnd

	serverPort := uint16(packet[pos])<<8 | uint16(packet[pos+1])
	pos += 2

	nextState, pos, ok, invalid := readMinecraftVarInt(packet, pos)
	if invalid || !ok || !isMinecraftNextState(nextState) || pos != len(packet) {
		return nil, false, false
	}

	return analyzer.PropMap{
		"protocol":        int(protocolVersion),
		"server_addr":     string(addrBytes),
		"server_port":     serverPort,
		"next_state":      int(nextState),
		"next_state_name": minecraftNextStateName(nextState),
		"packet_length":   int(packetLen),
	}, false, true
}

func readMinecraftVarInt(buf []byte, offset int) (value int32, next int, ok bool, invalid bool) {
	var result int32
	for i := 0; i < 5; i++ {
		if offset+i >= len(buf) {
			return 0, offset, false, false
		}
		b := buf[offset+i]
		result |= int32(b&0x7f) << (7 * i)
		if b&0x80 == 0 {
			return result, offset + i + 1, true, false
		}
	}
	return 0, offset, false, true
}

func isMinecraftNextState(state int32) bool {
	switch state {
	case minecraftNextStateStatus, minecraftNextStateLogin, minecraftNextStateTransfer:
		return true
	default:
		return false
	}
}

func minecraftNextStateName(state int32) string {
	switch state {
	case minecraftNextStateStatus:
		return "status"
	case minecraftNextStateLogin:
		return "login"
	case minecraftNextStateTransfer:
		return "transfer"
	default:
		return "unknown"
	}
}