package crypto

import (
	"crypto/aes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"fmt"

	"golang.org/x/crypto/curve25519"
)

type SharedSecret struct {
	aBytes [32]byte
}

func MakeSharedSecret(key []byte) SharedSecret {
	var secret SharedSecret
	copy(secret.aBytes[:], key)
	return secret
}

func MakeSharedSecretFromGroupSecret(group []byte) SharedSecret {
	var secret SharedSecret
	copy(secret.aBytes[:16], group)
	return secret
}

func (ss *SharedSecret) HMAC(message []byte) uint16 {
	h := hmac.New(sha256.New, ss.aBytes[:])
	h.Write(message)
	r := h.Sum(nil)
	return binary.BigEndian.Uint16(r[:2])
}

func (ss *SharedSecret) key() []byte {
	return ss.aBytes[:aes.BlockSize]
}

func (ss *SharedSecret) Decrypt(text []byte) ([]byte, error) {
	block, err := aes.NewCipher(ss.key())
	if err != nil {
		return nil, err
	}

	length := len(text)
	message := make([]byte, length)
	copy(message, text)

	const chunkSize = aes.BlockSize
	remain := length % chunkSize

	if remain > 0 {
		padding := chunkSize - remain
		message = append(message, make([]byte, padding)...)
	}

	for i, l := 0, len(message); i < l; i += chunkSize {
		block.Decrypt(message[i:i+chunkSize], message[i:i+chunkSize])
	}

	return message[:length], nil
}

func (ss *SharedSecret) MACThenDecrypt(text []byte, mac uint16) ([]byte, error) {
	if our := ss.HMAC(text); our != mac {
		return nil, fmt.Errorf("expected MAC %04X, got %04X", mac, our)
	}
	return ss.Decrypt(text)
}

func (ss *SharedSecret) Encrypt(message []byte) ([]byte, error) {
	block, err := aes.NewCipher(ss.key())
	if err != nil {
		return nil, err
	}

	text := make([]byte, len(message))
	copy(text, message)

	const chunkSize = aes.BlockSize
	remain := len(text) % chunkSize

	if remain > 0 {
		padding := chunkSize - remain
		text = append(text, make([]byte, padding)...)
	}

	for i, l := 0, len(text); i < l; i += chunkSize {
		block.Encrypt(text[i:i+chunkSize], text[i:i+chunkSize])
	}

	return text, nil
}

func (ss *SharedSecret) EncryptThenMAC(message []byte) (uint16, []byte, error) {
	text, err := ss.Encrypt(message)
	if err != nil {
		return 0, nil, err
	}
	return ss.HMAC(text), text, nil
}

type StaticSecret [32]byte

func (ss StaticSecret) PublicKey() (*PublicKey, error) {
	pub, err := curve25519.X25519(ss[:], curve25519.Basepoint)
	if err != nil {
		return nil, err
	}
	return NewPublicKey(pub)
}

func (ss StaticSecret) DiffieHellman(other *PublicKey) (SharedSecret, error) {
	shared, err := curve25519.X25519(ss[:], other.Bytes())
	if err != nil {
		return SharedSecret{}, err
	}

	var montgomery SharedSecret
	copy(montgomery.aBytes[:], shared)
	return montgomery, nil
}

type Signature [64]byte

func (sig Signature) MarshalJSON() ([]byte, error) {
	s := hex.EncodeToString(sig[:])
	return json.Marshal(s)
}

func (sig *Signature) UnmarshalJSON(b []byte) error {
	var s string
	if err := json.Unmarshal(b, &s); err != nil {
		return err
	}
	copy((*sig)[:], []byte(s))
	return nil
}