package main

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
)

const (
	HandshakePacketID = 0x00
	StatusRequestID   = 0x00
	StatusResponseID  = 0x00
	PingRequestID     = 0x01
	PongResponseID    = 0x01
)

type HandshakePacket struct {
	ProtocolVersion int
	ServerAddress   string
	ServerPort      uint16
	NextState       int
}

type StatusResponse struct {
	Version     VersionInfo `json:"version"`
	Players     PlayerInfo  `json:"players"`
	Description Chat        `json:"description"`
}

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

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

type Pong struct {
	Timestamp int64 `json:"timestamp"`
}

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

type Chat struct {
	Text string `json:"text"`
}

type ClientState struct {
	conn      net.Conn
	nextState int
}

func main() {
	listener, err := net.Listen("tcp", "localhost:25565")
	if err != nil {
		fmt.Println("Error starting server:", err)
		return
	}
	defer listener.Close()

	fmt.Println("Server started on port 25565")

	for {
		conn, err := listener.Accept()
		if err != nil {
			fmt.Println("Error accepting connection:", err)
			continue
		}

		client := &ClientState{conn: conn, nextState: 0}
		go handleConnection(client)
	}
}

func handleConnection(client *ClientState) {
	defer client.conn.Close()
	fmt.Println("New connection from:", client.conn.RemoteAddr())

	buf := make([]byte, 1024)
	for {
		n, err := client.conn.Read(buf)
		if err != nil {
			if err.Error() == "EOF" {
				fmt.Println("Client disconnected:", client.conn.RemoteAddr())
			} else {
				fmt.Println("Error reading from client:", err)
			}
			break
		}

		reader := bytes.NewReader(buf[:n])

		// Read packet length
		packetLength, err := readVarInt(reader)
		if err != nil {
			fmt.Println("Error reading packet length:", err)
			break
		}

		// Ensure full packet is read
		packetData := make([]byte, packetLength)
		if _, err := io.ReadFull(reader, packetData); err != nil {
			fmt.Println("Error reading full packet data:", err)
			break
		}

		packetReader := bytes.NewReader(packetData)

		packetID, err := readVarInt(packetReader)
		if err != nil {
			fmt.Println("Error reading packet ID:", err)
			break
		}

		fmt.Println("Packet ID:", packetID)

		switch packetID {
		case HandshakePacketID: // The same packet ID is used for the Status Request
			if client.nextState == 1 {
				handleStatusRequest(client)
			} else if client.nextState == 2 {
				fmt.Println("Received status request in the wrong state")
			} else if client.nextState == 3 {
				fmt.Println("Received status request in the wrong state")
			} else {
				handshake, err := handleHandshake(packetReader)
				if err != nil {
					fmt.Println("Error processing handshake:", err)
					return
				}
				client.nextState = handshake.NextState
			}
		case PingRequestID:
			var timestamp int64
			if err := binary.Read(packetReader, binary.BigEndian, &timestamp); err != nil {
				fmt.Println("Error reading timestamp:", err)
				break
			}

			var responsePacket bytes.Buffer
			writeVarInt(&responsePacket, 9) // Packet length: 1 byte for ID + 8 bytes for timestamp
			responsePacket.WriteByte(PongResponseID)
			binary.Write(&responsePacket, binary.BigEndian, timestamp)

			if _, err := client.conn.Write(responsePacket.Bytes()); err != nil {
				fmt.Println("Error sending pong response:", err)
			} else {
				fmt.Println("Pong response sent successfully:", timestamp)
			}
		}
	}
}

func handleHandshake(reader *bytes.Reader) (HandshakePacket, error) {
	var handshake HandshakePacket
	var err error

	handshake.ProtocolVersion, err = readVarInt(reader)
	if err != nil {
		return handshake, err
	}

	handshake.ServerAddress, err = readString(reader)
	if err != nil {
		return handshake, err
	}

	if err := binary.Read(reader, binary.BigEndian, &handshake.ServerPort); err != nil {
		return handshake, err
	}

	handshake.NextState, err = readVarInt(reader)
	if err != nil {
		return handshake, err
	}

	fmt.Printf("Handshake received: %+v\n", handshake)
	return handshake, nil
}

func handleStatusRequest(client *ClientState) {
	status := StatusResponse{
		Version:     VersionInfo{"1.21.4", 769},
		Players:     PlayerInfo{Max: 100, Online: 5, Sample: []Player{}},
		Description: Chat{Text: "A Minecraft Server"},
	}

	response, err := json.Marshal(status)
	if err != nil {
		fmt.Println("Error marshaling status response:", err)
		return
	}

	var packetData bytes.Buffer
	writeVarInt(&packetData, StatusResponseID)
	writeVarInt(&packetData, len(response))
	packetData.Write(response)

	var responsePacket bytes.Buffer
	writeVarInt(&responsePacket, packetData.Len())
	responsePacket.Write(packetData.Bytes())

	fmt.Println("Final Response Packet:", responsePacket.Bytes())

	if _, err = client.conn.Write(responsePacket.Bytes()); err != nil {
		fmt.Println("Error sending response:", err)
	} else {
		fmt.Println("Status response sent successfully")
	}
}

func writeVarInt(buf *bytes.Buffer, value int) {
	for {
		temp := byte(value & 0x7F)
		value >>= 7
		if value > 0 {
			temp |= 0x80
		}
		buf.WriteByte(temp)
		if value == 0 {
			break
		}
	}
}

func readVarInt(reader *bytes.Reader) (int, error) {
	value, position := 0, 0
	for {
		b, err := reader.ReadByte()
		if err != nil {
			return 0, err
		}
		value |= int(b&0x7F) << position
		if (b & 0x80) == 0 {
			break
		}
		position += 7
		if position >= 32 {
			return 0, errors.New("VarInt too big")
		}
	}
	return value, nil
}

func readString(reader *bytes.Reader) (string, error) {
	length, err := readVarInt(reader)
	if err != nil {
		return "", err
	}

	strBuf := make([]byte, length)
	if _, err := reader.Read(strBuf); err != nil {
		return "", err
	}
	return string(strBuf), nil
}