package crypto

import (
	"crypto/rand"
	"fmt"
	"lux/proto"
	"os"
)

type LuxKey struct {
	Type proto.LuxType // 2
	Id   proto.LuxID   // 16
	Key  []byte        // 32
	IV   []byte        // 16
}

const LUX_KEY_AES_SIZE = 32
const LUX_KEY_AES_BLOCKSIZE = 16
const LUX_KEY_IV_SIZE = LUX_KEY_AES_BLOCKSIZE
const LUX_PROTO_KEY_SIZE = proto.LUX_PROTO_TYPE_SIZE + proto.LUX_PROTO_ID_SIZE + LUX_KEY_AES_SIZE + LUX_KEY_IV_SIZE

func NewLuxKey(keyType proto.LuxType) (LuxKey, error) {
	key := LuxKey{
		Type: keyType,
		Id:   proto.NewLuxID(),
		Key:  make([]byte, LUX_KEY_AES_SIZE),
		IV:   make([]byte, LUX_KEY_IV_SIZE),
	}

	if _, err := rand.Read(key.Key); err != nil {
		return key, err
	}
	if _, err := rand.Read(key.IV); err != nil {
		return key, err
	}

	return key, nil
}

func (key *LuxKey) Read(rd *proto.LuxBuffer) error {
	var err error
	var _key []byte
	var iv []byte

	if err = key.Type.Read(rd); err != nil {
		return err
	}
	if err = key.Id.Read(rd); err != nil {
		return err
	}
	if _key, err = rd.ReadNext(LUX_KEY_AES_SIZE); err != nil {
		return err
	}
	if iv, err = rd.ReadNext(LUX_KEY_IV_SIZE); err != nil {
		return err
	}

	key.Key = _key
	key.IV = iv
	return nil
}

func (key *LuxKey) Write(wd *proto.LuxBuffer) {
	key.Type.Write(wd)
	key.Id.Write(wd)
	wd.WriteBytes(key.Key)
	wd.WriteBytes(key.IV)
}

func (key *LuxKey) String() string {
	return fmt.Sprintf("%s %s", key.Type.String(), key.Id.String())
}

const LUX_KEYSTORE_FILEMODE = os.FileMode(int(0600))

type LuxKeyStore struct {
	filePath string
	keys     map[proto.LuxID]LuxKey
}

func NewLuxKeyStore(filePath string) LuxKeyStore {
	return LuxKeyStore{
		filePath: filePath,
		keys:     make(map[proto.LuxID]LuxKey),
	}
}

func (ks *LuxKeyStore) Load() error {
	bytes, err := os.ReadFile(ks.filePath)
	if err != nil {
		// probably no file exists, so we create it
		return ks.Save()
	} else {
		rd := proto.FromSlice(bytes)
		keyNum := rd.Remaining() / LUX_PROTO_KEY_SIZE
		for i := 0; i < keyNum; i++ {
			var key LuxKey
			if err := key.Read(&rd); err != nil {
				return err
			}

			// add new key
			ks.keys[key.Id] = key
		}
	}

	return nil
}

func (ks *LuxKeyStore) Save() error {
	wd := proto.AllocLuxBuffer(len(ks.keys))
	for _, key := range ks.keys {
		key.Write(&wd)
	}

	return os.WriteFile(ks.filePath, wd.AllBytes(), LUX_KEYSTORE_FILEMODE)
}

func (ks *LuxKeyStore) Get(id proto.LuxID) (LuxKey, bool) {
	key, ok := ks.keys[id]
	return key, ok
}

func (ks *LuxKeyStore) Put(key LuxKey) error {
	ks.keys[key.Id] = key
	if err := ks.Save(); err != nil {
		return err
	}

	return nil
}

func (ks *LuxKeyStore) Count() int {
	return len(ks.keys)
}

func (ks *LuxKeyStore) Keys() []LuxKey {
	values := make([]LuxKey, len(ks.keys))

	i := 0
	for _, value := range ks.keys {
		values[i] = value
		i += 1
	}

	return values
}