package main

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"math/rand"
	"net"
	"strconv"
	"strings"
	"time"
)

type ServerStatus struct {
	Online      bool
	Players     []string
	OnlineCount int
	MaxPlayers  int
}

type slpVersion struct {
	Name     string `json:"name"`
	Protocol int    `json:"protocol"`
}

type slpSample struct {
	Name string `json:"name"`
	ID   string `json:"id"`
}

type slpPlayers struct {
	Max    int         `json:"max"`
	Online int         `json:"online"`
	Sample []slpSample `json:"sample"`
}

type slpResponse struct {
	Version     slpVersion `json:"version"`
	Players     slpPlayers `json:"players"`
	Description any        `json:"description"`
}

func queryServer(host string, port int) (*ServerStatus, error) {
	status, err := querySLP(host, port)
	if err != nil {
		return nil, err
	}

	// If SLP returned no player names, try GS4 query
	if len(status.Players) == 0 && status.OnlineCount > 0 {
		players, gs4Err := queryGS4(host, port)
		if gs4Err == nil && len(players) > 0 {
			status.Players = players
		}
	}

	return status, nil
}

func querySLP(host string, port int) (*ServerStatus, error) {
	addr := fmt.Sprintf("%s:%d", host, port)
	conn, err := net.DialTimeout("tcp", addr, 10*time.Second)
	if err != nil {
		return nil, fmt.Errorf("connect: %w", err)
	}
	defer conn.Close()

	if err := conn.SetDeadline(time.Now().Add(10 * time.Second)); err != nil {
		return nil, fmt.Errorf("set deadline: %w", err)
	}

	// Handshake
	{
		var payload strings.Builder
		writeVarInt(&payload, 767)     // protocol version
		writeString(&payload, host)    // server address
		_ = binary.Write(&payload, binary.BigEndian, uint16(port)) // server port
		writeVarInt(&payload, 1)       // next state: status

		packet := buildPacket(0, payload.String())
		if _, err := conn.Write([]byte(packet)); err != nil {
			return nil, fmt.Errorf("handshake write: %w", err)
		}
	}

	// Status request
	{
		packet := buildPacket(0, "")
		if _, err := conn.Write([]byte(packet)); err != nil {
			return nil, fmt.Errorf("request write: %w", err)
		}
	}

	// Read response
	reader := bufio.NewReader(conn)
	data, err := readPacket(reader)
	if err != nil {
		return nil, fmt.Errorf("read response: %w", err)
	}

	r := strings.NewReader(string(data))
	packetID, err := readVarInt(r)
	if err != nil {
		return nil, fmt.Errorf("read packet id: %w", err)
	}
	if packetID != 0 {
		return nil, fmt.Errorf("unexpected packet id: %d", packetID)
	}

	jsonStr, err := readString(r)
	if err != nil {
		return nil, fmt.Errorf("read json: %w", err)
	}

	var resp slpResponse
	if err := json.Unmarshal([]byte(jsonStr), &resp); err != nil {
		return nil, fmt.Errorf("parse json: %w", err)
	}

	status := &ServerStatus{
		Online:      true,
		OnlineCount: resp.Players.Online,
		MaxPlayers:  resp.Players.Max,
	}

	for _, s := range resp.Players.Sample {
		status.Players = append(status.Players, s.Name)
	}

	return status, nil
}

func queryGS4(host string, port int) ([]string, error) {
	addr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		return nil, err
	}

	conn, err := net.DialUDP("udp", nil, addr)
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	sessionID := int32(rand.Int31n(0x7FFFFFFF))

	// Handshake: FE FD 09 <session_id>
	conn.SetDeadline(time.Now().Add(10 * time.Second))
	hs := []byte{0xFE, 0xFD, 0x09}
	hs = binary.BigEndian.AppendUint32(hs, uint32(sessionID))
	conn.Write(hs)

	resp := make([]byte, 64)
	n, err := conn.Read(resp)
	if err != nil {
		return nil, fmt.Errorf("gs4 handshake: %w", err)
	}
	data := resp[:n]

	tokenStr := strings.TrimRight(string(data[5:]), "\x00")
	challengeToken, err := strconv.Atoi(tokenStr)
	if err != nil {
		return nil, fmt.Errorf("parse challenge: %w", err)
	}

	// Full stat query: FE FD 00 <session_id> <challenge_token> 00 00 00 00
	conn.SetDeadline(time.Now().Add(10 * time.Second))
	q := []byte{0xFE, 0xFD, 0x00}
	q = binary.BigEndian.AppendUint32(q, uint32(sessionID))
	q = binary.BigEndian.AppendUint32(q, uint32(challengeToken))
	q = append(q, 0x00, 0x00, 0x00, 0x00)
	conn.Write(q)

	resp = make([]byte, 65536)
	n, err = conn.Read(resp)
	if err != nil {
		return nil, fmt.Errorf("gs4 query: %w", err)
	}
	data = resp[:n]

	// Skip: type(1) + sessionID(4) + padding(11) = 16 bytes
	content := data[16:]

	sepIdx := bytes.Index(content, []byte{0x00, 0x00, 0x01})
	if sepIdx < 0 {
		return nil, fmt.Errorf("gs4: player separator not found")
	}

	playerData := content[sepIdx+3:]
	playerData = bytes.TrimPrefix(playerData, []byte("player_\x00\x00"))
	playerData = bytes.TrimRight(playerData, "\x00")
	players := bytes.Split(playerData, []byte{0x00})

	var result []string
	for _, p := range players {
		if len(p) > 0 {
			result = append(result, string(p))
		}
	}

	return result, nil
}

func buildPacket(id int, payload string) string {
	var inner strings.Builder
	writeVarInt(&inner, id)
	inner.WriteString(payload)

	var buf strings.Builder
	writeVarInt(&buf, inner.Len())
	buf.WriteString(inner.String())
	return buf.String()
}

func readPacket(r *bufio.Reader) ([]byte, error) {
	length, err := readVarInt(r)
	if err != nil {
		return nil, fmt.Errorf("read length: %w", err)
	}
	if length < 0 || length > 65536 {
		return nil, fmt.Errorf("invalid packet length: %d", length)
	}
	buf := make([]byte, length)
	if _, err := io.ReadFull(r, buf); err != nil {
		return nil, fmt.Errorf("read payload: %w", err)
	}
	return buf, nil
}

func readVarInt(r interface{ ReadByte() (byte, error) }) (int, error) {
	result := 0
	shift := 0
	for {
		b, err := r.ReadByte()
		if err != nil {
			return 0, err
		}
		result |= int(b&0x7F) << shift
		shift += 7
		if shift > 35 {
			return 0, fmt.Errorf("varint too long")
		}
		if b&0x80 == 0 {
			break
		}
	}
	return result, nil
}

func writeVarInt(w interface{ WriteByte(byte) error }, value int) {
	for value > 0x7F {
		w.WriteByte(byte(value&0x7F) | 0x80)
		value >>= 7
	}
	w.WriteByte(byte(value & 0x7F))
}

func readString(r interface{ ReadByte() (byte, error) }) (string, error) {
	length, err := readVarInt(r)
	if err != nil {
		return "", err
	}
	if length < 0 || length > 65536 {
		return "", fmt.Errorf("invalid string length: %d", length)
	}
	buf := make([]byte, length)
	if _, err := io.ReadFull(r.(io.Reader), buf); err != nil {
		return "", fmt.Errorf("read string: %w", err)
	}
	return string(buf), nil
}

func writeString(w interface{ WriteByte(byte) error }, s string) {
	writeVarInt(w, len(s))
	for i := 0; i < len(s); i++ {
		w.WriteByte(s[i])
	}
}