package crypto

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"lux/proto"
	"lux/rpc"
	"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 NewLuxTempKeyStore() LuxKeyStore {
	return NewLuxKeyStore("")
}

func (ks *LuxKeyStore) Load() error {
	if ks.filePath == "" {
		// skip for temp ks
		return nil
	}

	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 {
	if ks.filePath == "" {
		// skip for temp ks
		return nil
	}

	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
}

// RPC
func (ks *LuxKeyStore) GetRpcName() string {
	return "ks"
}

func (ks *LuxKeyStore) Register(sv *rpc.LuxRpcServer) {}

func LuxKeyStoreIntoRpc(ks *LuxKeyStore) rpc.LuxRpcKeyStore {
	keyStore := rpc.LuxRpcKeyStore{
		Hosts: make([]rpc.LuxRpcKeyHost, 0),
		Nodes: make([]rpc.LuxRpcKeyNode, 0),
	}

	for id, key := range ks.keys {
		if key.Type == proto.LuxTypeNode {
			keyStore.Nodes = append(keyStore.Nodes, rpc.LuxRpcKeyNode{
				ID:      id.String(),
				KeyBlob: base64.StdEncoding.EncodeToString(key.Key),
				IVBlob:  base64.StdEncoding.EncodeToString(key.IV),
			})
		} else {
			keyStore.Hosts = append(keyStore.Hosts, rpc.LuxRpcKeyHost{
				ID:      id.String(),
				KeyBlob: base64.StdEncoding.EncodeToString(key.Key),
				IVBlob:  base64.StdEncoding.EncodeToString(key.IV),
			})
		}
	}

	return keyStore
}

func LuxKeyStoreFromRpc(rpcKs rpc.LuxRpcKeyStore, savePath string) (LuxKeyStore, error) {
	ks := NewLuxKeyStore(savePath)

	for _, rpcKey := range rpcKs.Nodes {
		id, err := proto.ParseLuxID(rpcKey.ID)
		if err != nil {
			return ks, err
		}

		key, err := base64.StdEncoding.DecodeString(rpcKey.KeyBlob)
		if err != nil {
			return ks, err
		}

		iv, err := base64.StdEncoding.DecodeString(rpcKey.IVBlob)
		if err != nil {
			return ks, err
		}

		ks.Put(LuxKey{
			Type: proto.LuxTypeNode,
			Id:   id,
			Key:  key,
			IV:   iv,
		})
	}

	for _, rpcKey := range rpcKs.Hosts {
		id, err := proto.ParseLuxID(rpcKey.ID)
		if err != nil {
			return ks, err
		}

		key, err := base64.StdEncoding.DecodeString(rpcKey.KeyBlob)
		if err != nil {
			return ks, err
		}

		iv, err := base64.StdEncoding.DecodeString(rpcKey.IVBlob)
		if err != nil {
			return ks, err
		}

		ks.Put(LuxKey{
			Type: proto.LuxTypeHost,
			Id:   id,
			Key:  key,
			IV:   iv,
		})
	}

	if err := ks.Save(); err != nil {
		return ks, err
	}
	return ks, nil
}

func (ks *LuxKeyStore) Handle(request rpc.LuxRpcRequest, rpcType rpc.LuxRpcType) (rpc.LuxRpcResponse, rpc.LuxRpcError, bool) {
	var rpcRes rpc.LuxRpcResponse

	// only root can manage keystore
	if rpcType != rpc.LuxRpcTypeRoot {
		return rpcRes, rpc.LUX_RPC_ERROR_ACCESS_DENIED, false
	}

	if request.Command == "get" {
		// get all keys, so we need to construct xml object for that
		return rpc.LuxRpcResponse{
			Keystore: LuxKeyStoreIntoRpc(ks),
		}, rpc.LuxRpcError{}, true
	}

	return rpcRes, rpc.LUX_RPC_ERROR_UNKNOWN_COMMAND, false
}