package net

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"errors"
	"log"
	"lux/crypto"
	"lux/proto"
	"net"
)

type LuxPacketType uint

const (
	LuxPacketTypeHeartbeat = 0
	LuxPacketTypeSync      = 1
)

const LUX_PROTO_PACKET_HDRLEN = proto.LUX_PROTO_ID_SIZE + 2

// Target, Type are decoded from header
// Buffer contains LuxBuffer with stripped headers
// NOTE: Packets are always padded, so payloads may be
// bigger than expected and contain trailing bytes

// Target (ID) is rather confusing. What it actually describes is
// associated key with route.
type LuxPacket struct {
	Target      proto.LuxID
	Type        LuxPacketType
	Buffer      proto.LuxBuffer
	ChannelType LuxChannelType
}

func DecryptLuxPacket(dgram LuxDatagram, key crypto.LuxKey) (LuxPacket, error) {
	packet := LuxPacket{}

	cipherBlock, err := aes.NewCipher(key.Key)
	if err != nil {
		return packet, err
	}
	decrypter := cipher.NewCBCDecrypter(cipherBlock, key.IV)

	payloadLen := len(dgram.Payload)
	if payloadLen == 0 || payloadLen%cipherBlock.BlockSize() != 0 {
		return packet, errors.New("payload is not block aligned")
	}

	payload := make([]byte, payloadLen)
	decrypter.CryptBlocks(payload, dgram.Payload)

	packet.Buffer = proto.FromSlice(payload)
	if err := packet.Target.Read(&packet.Buffer); err != nil {
		return packet, err

	}

	if pType, err := packet.Buffer.ReadUint16(); err == nil {
		packet.Type = LuxPacketType(pType)
	} else {
		return packet, err
	}

	data, err := packet.Buffer.ReadNext(packet.Buffer.Remaining())
	if err != nil {
		return packet, err
	}

	// strip lux packet headers from payload, as well as padding
	packet.Buffer = proto.FromSlice(data)
	return packet, nil
}

func EncryptLuxPacket(packet LuxPacket, key crypto.LuxKey, target *net.UDPAddr) (LuxDatagram, error) {
	dgram := LuxDatagram{}

	cipherBlock, err := aes.NewCipher(key.Key)
	if err != nil {
		return dgram, err
	}
	encrypter := cipher.NewCBCEncrypter(cipherBlock, key.IV)

	var paddingLen int
	packetLen := LUX_PROTO_PACKET_HDRLEN + packet.Buffer.Length()
	if packetLen%encrypter.BlockSize() != 0 {
		//paddingLen = ((packetLen / encrypter.BlockSize()) + 1) * encrypter.BlockSize()
		paddingLen = encrypter.BlockSize() - (packetLen % encrypter.BlockSize())
		log.Default().Print(packetLen, paddingLen)
	} else {
		paddingLen = 0
	}

	padding := make([]byte, paddingLen)
	_, err = rand.Read(padding)
	if err != nil {
		return dgram, err
	}

	wd := proto.AllocLuxBuffer(packetLen + paddingLen)
	key.Id.Write(&wd)
	wd.WriteUint16(uint16(packet.Type))
	wd.WriteBytes(packet.Buffer.AllBytes())
	wd.WriteBytes(padding)

	dgram.Target = target
	dgram.Payload = make([]byte, wd.Length())

	encrypter.CryptBlocks(dgram.Payload, wd.AllBytes())
	return dgram, nil
}