package protocol

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"fmt"
)

// Crypto Globals for Handshake (as per spec)
var (
	HandshakeMac   = [8]byte{0x54, 0x53, 0x33, 0x49, 0x4E, 0x49, 0x54, 0x31}                                                // "TS3INIT1"
	HandshakeKey   = []byte{0x63, 0x3A, 0x5C, 0x77, 0x69, 0x6E, 0x64, 0x6F, 0x77, 0x73, 0x5C, 0x73, 0x79, 0x73, 0x74, 0x65} // "c:\windows\syste"
	HandshakeNonce = []byte{0x6D, 0x5C, 0x66, 0x69, 0x72, 0x65, 0x77, 0x61, 0x6C, 0x6C, 0x33, 0x32, 0x2E, 0x63, 0x70, 0x6C} // "m\firewall32.cpl"
)

// EAX Mode Implementation (AES-CTR + OMAC) -> For simplicity we will use a simplified approach or find a library if possible.
// Since Go doesn't have EAX in standard lib, and we want to avoid complex dependencies if possible for this snippet,
// we'll implement the key generation logic. The EAX encryption usually wraps AES-CTR.
// Important: For a production client, use a verified crypto library.

type CryptoState struct {
	SharedIV     []byte
	SharedMac    []byte
	GenerationID uint32
}

// GenerateKeyNonce generates the Key and Nonce for EAX encryption/decryption as per section 1.6.2
func (s *CryptoState) GenerateKeyNonce(header *PacketHeader, isClientToServer bool) ([]byte, []byte) {
	// 1. Temporary buffer
	// length depends on protocol version/logic.
	// Spec: "The old protocol SharedIV ... will be 20 bytes long ... new protocol SharedIV ... will have 64 bytes"
	// GenerationID starts at 0.

	var temp []byte
	sivLen := len(s.SharedIV)

	if sivLen == 20 {
		temp = make([]byte, 26)
	} else {
		// Spec 1.6.2 New Protocol: 1 (type) + 1 (pt) + 4 (pgid) + 64 (sharedIV) = 70 bytes.
		temp = make([]byte, 70)
	}

	// temp[0] 'c'/'s' logic:
	// Spec: "0x30 for Client, 0x31 for Server... Wait.
	// "Used for Client -> Server: 0x31"
	// "Used for Server -> Client: 0x30"
	if isClientToServer {
		temp[0] = 0x31
	} else { // Client <- Server
		temp[0] = 0x30
	}

	// temp[1] = PT (Base Type only, no flags!)
	// ts3j uses getType().getIndex() which is the enum value.
	temp[1] = header.Type & 0x0F

	// temp[2..6] = PGId (Network Order)
	binary.BigEndian.PutUint32(temp[2:6], s.GenerationID)

	// Copy SharedIV
	// if SIV.length == 20 -> temp[6..26] = SIV[0..20]
	// if SIV.length == 64 -> temp[6..70] = SIV[0..64]
	copy(temp[6:], s.SharedIV)

	// keynonce = sha256(temp)
	hash := sha256.Sum256(temp)

	key := make([]byte, 16)
	nonce := make([]byte, 16)

	copy(key, hash[0:16])
	copy(nonce, hash[16:32])

	// XOR Key with PID
	// key[0] = key[0] xor ((PId & 0xFF00) >> 8)
	// key[1] = key[1] xor ((PId & 0x00FF) >> 0)
	key[0] ^= byte((header.PacketID & 0xFF00) >> 8)
	key[1] ^= byte((header.PacketID & 0x00FF))

	return key, nonce
}

// Base64 Helpers for TS3
func Base64Encode(data []byte) string {
	return base64.StdEncoding.EncodeToString(data)
}

func Base64Decode(s string) ([]byte, error) {
	return base64.StdEncoding.DecodeString(s)
}

// SHA1 Helper
func Sha1(data []byte) []byte {
	h := sha1.New()
	h.Write(data)
	return h.Sum(nil)
}

// GenerateHashCash implements the puzzle solver (Section 4.1)
// Finds the level of a public key + offset
func GetHashCashLevel(key []byte, offset uint64) int {
	// data = sha1(publicKey + keyOffset)
	// Note: Reference implementation concatenates the string representation of offset?
	// Spec says: "The key offset is a u64 number, which gets converted to a string when concatenated."

	offsetStr := fmt.Sprintf("%d", offset)
	buf := append(key, []byte(offsetStr)...)
	hash := Sha1(buf)

	// Count leading zero bits
	level := 0
	for _, b := range hash {
		if b == 0 {
			level += 8
		} else {
			// Count bits in this byte
			for i := 7; i >= 0; i-- {
				if (b & (1 << i)) == 0 {
					level++
				} else {
					return level
				}
			}
			break
		}
	}
	return level
}

// Fake encryption for now to proceed until EAX is fully implemented or needed
// The spec requires EAX for commands.
func EncryptPacket(pkt *Packet, key, nonce []byte) error {
	// Placeholder: In a real implementation, we would use an EAX implementation here.
	// For the initial handshake (Init1 packets), encryption is often disabled or simplified (XOR).
	// Reviewing: Init1 packets are NOT encrypted (FlagUnencrypted set).
	// Command packets ARE encrypted.

	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}

	// Standard CTR implementation (part of EAX)
	// EAX components: CTR for encryption, OMAC for authentication (MAC)
	// We strictly need the MAC for the packet header.

	// For minimal functional requirement "Connect", we might need full EAX if the server enforces it on the first Command.
	// For now we will implement standard CTR for the data payload.
	// The MAC calculation is complex and requires OMAC1.

	stream := cipher.NewCTR(block, nonce)
	stream.XORKeyStream(pkt.Data, pkt.Data)

	return nil
}

func DecryptPacket(pkt *Packet, key, nonce []byte) error {
	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}

	stream := cipher.NewCTR(block, nonce)
	stream.XORKeyStream(pkt.Data, pkt.Data)

	return nil
}

// Curve25519 Helper
func GenerateECDHKeypair() (public, private [32]byte, err error) {
	// Generate private key (use real random in prod)
	// For reproduction/dev we might use fixed.
	// Actually we need `crypto/rand`
	// private = rand(32)
	// curve25519.ScalarBaseMult(&public, &private)
	return
}