package main

import (
	"bufio"
	"bytes"
	"compress/zlib"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/x509"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/google/uuid"
	"io"
	"log"
	"math"
	"net/http"
	"regexp"
	"strings"
)

func generateRandomBytes(size int) ([]byte, error) {
	randomBytes := make([]byte, size)
	_, err := rand.Read(randomBytes)
	if err != nil {
		return nil, err
	}
	return randomBytes, nil
}

type TextPiece struct {
	Text      string `json:"text,omitempty"`
	Color     string `json:"color,omitempty"`
	CleanText string `json:"cleantext,omitempty"`
}

type TextComponent struct {
	Text      string      `json:"text,omitempty"`
	Extra     []TextPiece `json:"extra,omitempty"`
	CleanText string      `json:"cleantext,omitempty"`
}

func getTextComponent(text string) TextComponent {
	return TextComponent{
		Text:      text,
		Extra:     nil,
		CleanText: text,
	}
}

type RegistryEntry struct {
	EntryID string `json:"entryID,omitempty"`
	HasNBT  bool   `json:"hasNBT,omitempty"`
	//TODO implement the nbt data
	NBTData *bytes.Buffer `json:"nbtData,omitempty"`
}

type RegistryData struct {
	RegistryID string          `json:"registryID,omitempty"`
	Entries    []RegistryEntry `json:"entries,omitempty"`
}

type TagArray struct {
	TagName string  `json:"tagName,omitempty"`
	Entries []int32 `json:"entries,omitempty"`
}

type UpdateTag struct {
	TagRegistryIdentifier string     `json:"tagRegistryIdentifier,omitempty"`
	Tags                  []TagArray `json:"tags,omitempty"`
}

type DatapackInfo struct {
	Namespace string `json:"namespace,omitempty"`
	ID        string `json:"id,omitempty"`
	Version   string `json:"version,omitempty"`
}

type ServerLink struct {
	LabelString string `json:"labelString,omitempty"`
	LabelInt    int32  `json:"labelInt,omitempty"`
	URL         string `json:"url,omitempty"`
}

// GenerateRSAKeyPair generates an RSA key pair and returns the public and private keys.
func GenerateRSAKeyPair() (*rsa.PrivateKey, *rsa.PublicKey, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, 1024)
	if err != nil {
		return nil, nil, err
	}
	return privateKey, &privateKey.PublicKey, nil
}

// ExportRSAPublicKey exports the public key in ASN.1 DER format.
func ExportRSAPublicKey(pubKey *rsa.PublicKey) ([]byte, error) {
	// Marshal the public key to ASN.1 DER format
	pubASN1, err := x509.MarshalPKIXPublicKey(pubKey)
	if err != nil {
		return nil, err
	}
	// Return the raw DER-encoded bytes
	return pubASN1, nil
}

type cfb8 struct {
	b         cipher.Block
	blockSize int
	in        []byte
	out       []byte

	decrypt bool
}

func (x *cfb8) XORKeyStream(dst, src []byte) {
	for i := range src {
		x.b.Encrypt(x.out, x.in)
		copy(x.in[:x.blockSize-1], x.in[1:])
		if x.decrypt {
			x.in[x.blockSize-1] = src[i]
		}
		dst[i] = src[i] ^ x.out[0]
		if !x.decrypt {
			x.in[x.blockSize-1] = dst[i]
		}
	}
}

// NewCFB8Encryptor returns a Stream which encrypts with cipher feedback mode
// (segment size = 8), using the given Block. The iv must be the same length as
// the Block's block size.
func newCFB8Encryptor(block cipher.Block, iv []byte) cipher.Stream {
	return newCFB8(block, iv, false)
}

// NewCFB8Decryptor returns a Stream which decrypts with cipher feedback mode
// (segment size = 8), using the given Block. The iv must be the same length as
// the Block's block size.
func newCFB8Decryptor(block cipher.Block, iv []byte) cipher.Stream {
	return newCFB8(block, iv, true)
}

func newCFB8(block cipher.Block, iv []byte, decrypt bool) cipher.Stream {
	blockSize := block.BlockSize()
	if len(iv) != blockSize {
		// stack trace will indicate whether it was de or encryption
		panic("cipher.newCFB: IV length must equal block size")
	}
	x := &cfb8{
		b:         block,
		blockSize: blockSize,
		out:       make([]byte, blockSize),
		in:        make([]byte, blockSize),
		decrypt:   decrypt,
	}
	copy(x.in, iv)

	return x
}

// NewCFB8Cipher Initialize AES/CFB8 encryption using the shared secret
func NewCFB8Cipher(secret []byte) (cipher.Stream, cipher.Stream, error) {
	block, err := aes.NewCipher(secret)
	if err != nil {
		return nil, nil, err
	}

	// Create CFB8 encryption/decryption ciphers using the custom implementation
	iv := secret // Initial vector (IV) is the same as the secret in this case
	encryptStream := newCFB8Encryptor(block, iv)
	decryptStream := newCFB8Decryptor(block, iv)

	return encryptStream, decryptStream, nil
}

func AuthDigest(parts [][]byte) string {
	h := sha1.New()

	// Hash each part incrementally
	for _, part := range parts {
		io.WriteString(h, string(part))
	}

	// Finalize the hash
	hash := h.Sum(nil)

	// Check for negative hashes
	negative := (hash[0] & 0x80) == 0x80
	if negative {
		hash = twosComplement(hash)
	}

	// Trim away zeroes
	res := strings.TrimLeft(hex.EncodeToString(hash), "0")
	if negative {
		res = "-" + res
	}

	return res
}

// little endian
func twosComplement(p []byte) []byte {
	carry := true
	for i := len(p) - 1; i >= 0; i-- {
		p[i] = ^p[i]
		if carry {
			carry = p[i] == 0xff
			p[i]++
		}
	}
	return p
}

type MojangProfile struct {
	ID         string           `json:"id"`
	Name       string           `json:"name"`
	Properties []MojangProperty `json:"properties"`
}

type MojangProperty struct {
	Name      string `json:"name"`
	Value     string `json:"value"`
	Signature string `json:"signature"`
}

func (mojangProperty *MojangProperty) add(buf *bytes.Buffer) {
	addString(buf, mojangProperty.Name)
	addString(buf, mojangProperty.Value)
	hasSignature := mojangProperty.Signature != ""
	addBool(buf, hasSignature)
	if hasSignature {
		addString(buf, mojangProperty.Signature)
	}
}

type PlayerAction struct {
	ActionMask     byte
	Name           string
	Properties     []MojangProperty
	GameMode       int32
	Listed         bool
	Ping           int32
	DisplayName    string
	HasDisplayName bool
}

// Helper to add a player with properties
func (playerAction *PlayerAction) add(buf *bytes.Buffer) {
	addString(buf, playerAction.Name)

	// Add properties
	addVarint(buf, int32(len(playerAction.Properties)))
	for _, prop := range playerAction.Properties {
		prop.add(buf)
	}
}

// Helper to update game mode

// Helper to update listed state

// Helper to update latency (ping)

// Helper to update display name

func (player *Player) hasJoinedSession(sharedSecret []byte) (bool, error) {
	serverKey, err := ExportRSAPublicKey(serverPublicKey)
	serverIDHash := AuthDigest([][]byte{serverID, sharedSecret, serverKey})
	remoteAddr := strings.Split(player.conn.RemoteAddr().String(), ":")[0]
	url := fmt.Sprintf("https://sessionserver.mojang.com/session/minecraft/hasJoined?username=%s&serverId=%s&ip=%s", player.Entity.Name, serverIDHash, remoteAddr)

	resp, err := http.Get(url)
	if err != nil {
		return false, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return false, fmt.Errorf("failed to verify session: %s", resp.Status)
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return false, err
	}

	var profile MojangProfile
	err = json.Unmarshal(body, &profile)
	if err != nil {
		return false, err
	}
	player.profile = profile
	return true, nil
}

// EncryptPacket encrypts a packet using the AES/CFB8 cipher stream.
func EncryptPacket(stream cipher.Stream, data *bytes.Buffer) *bytes.Buffer {
	if stream == nil {
		return data
	}
	encryptedData := make([]byte, data.Len())
	stream.XORKeyStream(encryptedData, data.Bytes())
	return bytes.NewBuffer(encryptedData)
}

// DecryptingReader returns a *bufio.Reader that provides decrypted data from the original reader.
func DecryptingReader(stream cipher.Stream, reader *bufio.Reader) *bufio.Reader {
	pr, pw := io.Pipe() // create a pipe to stream decrypted data

	go func() {
		defer pw.Close()          // ensure the pipe writer is closed when done
		buf := make([]byte, 4096) // buffer size, adjust as needed
		for {
			n, err := reader.Read(buf)
			if err != nil {
				if err == io.EOF {
					break
				}
				pw.CloseWithError(err) // propagate any non-EOF error to the pipe reader
				return
			}

			data := buf[:n]
			decryptedBlock := make([]byte, len(data))
			if stream != nil {
				stream.XORKeyStream(decryptedBlock, data) // decrypt the data
			} else {
				decryptedBlock = data // if no stream, just pass the data through
			}

			_, err = pw.Write(decryptedBlock) // write decrypted data to the pipe writer
			if err != nil {
				pw.CloseWithError(err) // propagate write errors
				return
			}
		}
	}()

	return bufio.NewReader(pr) // wrap the pipe reader in a bufio.Reader
}

// Function to clean Minecraft formatting from a text
func cleanMinecraftFormatting(text string) string {
	// Regex to match Minecraft formatting codes like § and &
	regex := regexp.MustCompile(`(?i)§[0-9A-FK-OR]`)
	return regex.ReplaceAllString(text, "")
}

// Decompress function using zlib
func decompress(data *bytes.Buffer) (buf *bytes.Buffer, err error) {
	reader, err := zlib.NewReader(data)
	if err != nil {
		return
	}
	defer func(reader io.ReadCloser) {
		_ = reader.Close()
	}(reader)

	var decompressedData bytes.Buffer
	_, err = io.Copy(&decompressedData, reader)
	if err != nil {
		return
	}

	return &decompressedData, nil
}

// Compress function using zlib
func compress(data *bytes.Buffer) (*bytes.Buffer, error) {
	var compressedData bytes.Buffer
	writer := zlib.NewWriter(&compressedData)

	_, err := writer.Write(data.Bytes())
	if err != nil {
		return nil, err
	}

	err = writer.Close()
	if err != nil {
		return nil, err
	}

	return &compressedData, nil
}

// Modified version of `sendPacket` with hexdump logging, including unencrypted data
func (player *Player) sendPacket(packetID int32, packetData *bytes.Buffer) (err error) {
	var origDataBuffer bytes.Buffer
	addVarint(&origDataBuffer, packetID) // Add the Packet ID as a VarInt
	_, err = origDataBuffer.ReadFrom(packetData)
	if err != nil {
		return err
	}

	var outBuffer bytes.Buffer
	var outTempBuffer bytes.Buffer
	length := int32(origDataBuffer.Len()) // Get the length of the uncompressed packet
	compressed := false                   // Track if the packet is compressed

	// Compression based on a threshold
	if player.compressionThreshold >= 0 && length >= player.compressionThreshold {
		compressed = true // Mark that compression was applied
		compressedData, err := compress(&origDataBuffer)
		if err != nil {
			return err
		}
		origDataBuffer = *compressedData

		// Add the uncompressed length to outTempBuffer
		addVarint(&outTempBuffer, length)
	}
	encrypted := player.encryptStream != nil

	// Log the unencrypted packet data
	log.Printf("Unencrypted Packet ID: %d", packetID)
	log.Printf("Unencrypted Packet Length: %d (compressed: %v, encrypted: %v)", origDataBuffer.Len(), compressed, encrypted)
	log.Printf("Hexdump of Unencrypted Packet:\n%s", hex.Dump(origDataBuffer.Bytes()))

	// Encrypt the (compressed or uncompressed) data buffer
	dataBuffer := origDataBuffer

	// Append the (compressed or uncompressed) data buffer to outTempBuffer
	_, err = outTempBuffer.ReadFrom(&dataBuffer)
	if err != nil {
		return err
	}

	// Add the total packet length to outBuffer
	addVarint(&outBuffer, int32(outTempBuffer.Len()))
	_, err = outBuffer.ReadFrom(&outTempBuffer)
	if err != nil {
		return err
	}

	if encrypted {
		outBuffer = *EncryptPacket(player.encryptStream, &outBuffer)
	}
	_, err = player.conn.Write(outBuffer.Bytes())
	if err != nil {
		return err
	}
	return nil
}

// Modified version of `readPacket` with hexdump logging
func (player *Player) readPacket() (packetID int32, packetData *bytes.Buffer, err error) {
	packetData = &bytes.Buffer{}
	finalReader := player.reader
	if player.decryptStream != nil {
		finalReader = DecryptingReader(player.decryptStream, player.reader)
	}
	packetLength, _, err := readVarint(finalReader)
	if err != nil {
		return
	}

	packetDataBytes := make([]byte, packetLength)
	nt, err := io.ReadFull(finalReader, packetDataBytes)
	if err != nil {
		return
	}

	_, err = packetData.Write(packetDataBytes)
	if err != nil {
		return 0, &bytes.Buffer{}, err
	}
	n := int32(nt)
	if n != packetLength {
		err = errors.New("packet read length mismatch")
		return
	}

	var dataLength int
	compressed := false // Track if the packet is compressed
	if player.compressionThreshold > 0 {
		dataLengthTemp, errTemp := receiveVarint(packetData)
		err = errTemp
		dataLength = int(dataLengthTemp)
	}

	// If dataLength > 0, it means the packet is compressed
	if dataLength > 0 {
		compressed = true
		packetData, err = decompress(packetData)
		if err != nil {
			return
		}
	}

	// Extract packet ID
	packetID, err = receiveVarint(packetData)

	// Log the packet information
	log.Printf("Received Packet ID: %d", packetID)
	log.Printf("Packet Length: %d (compressed: %v)", packetLength, compressed)
	log.Printf("Hexdump of Decoded Packet:\n%s", hex.Dump(packetData.Bytes()))

	return
}

func addVarint(buffer *bytes.Buffer, value int32) {
	uValue := uint32(value) // Explicitly cast to an unsigned type
	for {
		b := byte(uValue & 0x7F)
		uValue >>= 7

		if (uValue != 0 && uValue != 0xFFFFFFFF) || (b&0x40 != 0) {
			b |= 0x80
		}

		addByte(buffer, b)

		if uValue == 0 || uValue == 0xFFFFFFFF {
			break
		}
	}
}

func addVarlong(buffer *bytes.Buffer, value int64) {
	uValue := uint64(value) // Explicitly cast to an unsigned type
	for {
		b := byte(uValue & 0x7F)
		uValue >>= 7

		if (uValue != 0 && uValue != 0xFFFFFFFFFFFFFFFF) || (b&0x40 != 0) {
			b |= 0x80
		}

		addByte(buffer, b)

		if uValue == 0 || uValue == 0xFFFFFFFFFFFFFFFF {
			break
		}
	}
}

func readVarint(r *bufio.Reader) (value int32, length uint32, err error) {
	var (
		b      byte
		shift  uint64
		uValue uint32
	)

	for {
		b, err = r.ReadByte()
		if err != nil {
			return 0, 0, err
		}
		length++

		uValue |= (uint32(b) & 0x7F) << shift
		shift += 7

		if b&0x80 == 0 {
			break
		}
	}

	if shift < 8*uint64(length) && b&0x40 != 0 {
		uValue |= ^uint32(0) << shift
	}
	value = int32(uValue)
	return
}

func receiveVarint(buf *bytes.Buffer) (value int32, err error) {
	var (
		b      byte
		shift  uint64
		uValue uint32
		length uint32
	)

	for i := 0; i < buf.Len(); i++ {
		b, err = buf.ReadByte()
		if err != nil {
			return
		}
		length++

		uValue |= (uint32(b) & 0x7F) << shift
		shift += 7

		if b&0x80 == 0 {
			break
		}
	}

	if shift < 8*uint64(length) && b&0x40 != 0 {
		uValue |= ^uint32(0) << shift
	}
	value = int32(uValue)
	return
}

func addUUID(buffer *bytes.Buffer, inUUID uuid.UUID) int32 {
	binaryUUID, err := inUUID.MarshalBinary()
	if err != nil {
		log.Println("Error marshalling UUID:", err)
		return 0
	}
	buffer.Write(binaryUUID)
	return int32(len(binaryUUID))
}

// Decodes a UUID from the buffer.
func receiveUUID(buffer *bytes.Buffer) (value uuid.UUID, err error) {
	b := make([]byte, 16)
	_, err = buffer.Read(b)
	if err != nil {
		return uuid.UUID{}, err
	}
	value, err = uuid.FromBytes(b)
	if err != nil {
		return uuid.UUID{}, err
	}
	return value, nil
}

func addByte(buffer *bytes.Buffer, byte byte) {
	buffer.WriteByte(byte)
}

func addBytes(buffer *bytes.Buffer, bytes []byte) {
	buffer.Write(bytes)
}

func addBool(buffer *bytes.Buffer, boolean bool) {
	var boolByte byte = 0x00
	if boolean {
		boolByte = 0x01
	}
	buffer.WriteByte(boolByte)
}

func receiveBool(buffer *bytes.Buffer) (boolean bool, err error) {
	b, err := buffer.ReadByte()
	if err != nil {
		return false, err
	}
	boolean = b == 1
	return
}

// Decodes an uint16 from the buffer.
func receiveUint16(buffer *bytes.Buffer) (value uint16) {
	// Create a byte slice to hold the 2 bytes for the uint16.
	b := make([]byte, 2)

	// Read exactly 2 bytes from the buffer.
	_, err := buffer.Read(b)
	if err != nil {
		return 0
	}

	// Convert the byte slice to uint16 using BigEndian.
	value = binary.BigEndian.Uint16(b)
	return value
}

// Decodes an uint16 from the buffer.
func receiveInt16(buffer *bytes.Buffer) (value int16) {
	// Create a byte slice to hold the 2 bytes for the uint16.
	b := make([]byte, 2)

	// Read exactly 2 bytes from the buffer.
	_, err := buffer.Read(b)
	if err != nil {
		return 0
	}

	// Convert the byte slice to uint16 using BigEndian.
	value = int16(binary.BigEndian.Uint16(b))
	return value
}

func addInt64(buffer *bytes.Buffer, value int64) {
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, uint64(value))
	buffer.Write(b)
}

func receiveInt64(buffer *bytes.Buffer) (value int64) {
	b := make([]byte, 8)
	_, err := buffer.Read(b)
	if err != nil {
		return 0
	}
	value = int64(binary.BigEndian.Uint64(b))
	return
}

func addFloat32(buffer *bytes.Buffer, value float32) {
	bits := math.Float32bits(value)
	b := make([]byte, 4)
	binary.BigEndian.PutUint32(b, bits)
	buffer.Write(b)
}

// Decodes a float32 from the buffer.
func receiveFloat32(buffer *bytes.Buffer) (value float32) {
	b := make([]byte, 4)
	_, err := buffer.Read(b)
	if err != nil {
		return 0
	}
	value = math.Float32frombits(binary.BigEndian.Uint32(b))
	return
}

func addFloat64(buffer *bytes.Buffer, value float64) {
	bits := math.Float64bits(value)
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, bits)
	buffer.Write(b)
}

// Decodes a float64 from the buffer.

func addInt32(buffer *bytes.Buffer, value int32) {
	b := make([]byte, 4)
	binary.BigEndian.PutUint32(b, uint32(value))
	buffer.Write(b)
}

// Decodes an int32 from the buffer.

func addInt16(buffer *bytes.Buffer, value int16) {
	b := make([]byte, 2)
	binary.BigEndian.PutUint16(b, uint16(value))
	buffer.Write(b)
}

func addUint32(buffer *bytes.Buffer, value uint32) {
	b := make([]byte, 4)
	binary.BigEndian.PutUint32(b, value)
	buffer.Write(b)
}

// Decodes an uint32 from the buffer.
func receiveUint32(buffer *bytes.Buffer) (value uint32, err error) {
	b := make([]byte, 4)
	if _, err = buffer.Read(b); err != nil {
		return 0, err
	}
	value = binary.BigEndian.Uint32(b)
	return
}

func addString(buffer *bytes.Buffer, str string) {
	// Encode string length using VarInt.
	addVarint(buffer, int32(len(str)))
	// Append the string bytes.
	buffer.Write([]byte(str))
}

// Decodes a string from the buffer.
func receiveString(buf *bytes.Buffer) (strOut string, err error) {
	stringLen, err := receiveVarint(buf)
	if err != nil {
		return "", err
	}
	stringBytes := make([]byte, stringLen)
	_, err = buf.Read(stringBytes)
	if err != nil {
		return "", err
	}
	strOut = string(stringBytes)
	return
}

// addTextComponent serializes a TextComponent to JSON and adds it to the buffer
func addTextComponent(buffer *bytes.Buffer, component TextComponent) {
	// Serialize the TextComponent to a JSON string
	jsonBytes, err := json.Marshal(component)
	if err != nil {
		return
	}

	// Convert JSON bytes to a string
	jsonString := string(jsonBytes)

	// Use addString to addExplosion the serialized TextComponent JSON to the buffer
	addString(buffer, jsonString)
	return
}