// Cryptographic key types, generation, loading, and box encryption helpers. Uses Ed25519
// for signing and NaCl box (Curve25519) for payload encryption.
package cryptography

import (
	"crypto/rand"
	"fmt"
	"os"
	"path/filepath"

	"github.com/libp2p/go-libp2p/core/crypto"
	"github.com/libp2p/go-libp2p/core/peer"
	"golang.org/x/crypto/curve25519"
	"golang.org/x/crypto/nacl/box"
)

// KeyPair holds an Ed25519 private/public keypair used for libp2p identity and signing.
type KeyPair struct {
	PrivateKey crypto.PrivKey
	PublicKey  crypto.PubKey
}

// BoxKeyPair holds a NaCl box (Curve25519) keypair for anonymous sealed-box encryption.
type BoxKeyPair struct {
	PrivateKey *[32]byte
	PublicKey  *[32]byte
}

// OperatorKey bundles the Ed25519 identity, NaCl box keys for message decryption,
// and a 32-byte auth token shared with implants at build time.
type OperatorKey struct {
	KeyPair
	PeerID    peer.ID
	BoxKeys   *BoxKeyPair
	AuthToken []byte
}

// ImplantKey bundles the Ed25519 identity and stores the operator's public key for
// verifying incoming command signatures.
type ImplantKey struct {
	KeyPair
	PeerID          peer.ID
	OperatorPubKey crypto.PubKey
}

// GenerateBoxKey creates a fresh NaCl box keypair for anonymous sealed-box encryption.
func GenerateBoxKey() (*BoxKeyPair, error) {
	pub, priv, err := box.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("generate box key: %w", err)
	}
	return &BoxKeyPair{PrivateKey: priv, PublicKey: pub}, nil
}

// LoadBoxKey derives a box keypair from a 32-byte private key seed.
func LoadBoxKey(data []byte) *BoxKeyPair {
	var priv [32]byte
	copy(priv[:], data)
	pub, err := curve25519.X25519(priv[:], curve25519.Basepoint)
	if err != nil {
		return nil
	}
	var pubArr [32]byte
	copy(pubArr[:], pub)
	return &BoxKeyPair{PrivateKey: &priv, PublicKey: &pubArr}
}

// Marshal returns the box private key as 32 raw bytes.
func (kp *BoxKeyPair) Marshal() []byte {
	return kp.PrivateKey[:]
}

// GenerateOperatorKey creates a new Ed25519 identity, box keypair, and 32-byte auth token.
func GenerateOperatorKey() (*OperatorKey, error) {
	priv, pub, err := crypto.GenerateEd25519Key(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("generate operator key: %w", err)
	}
	pid, err := peer.IDFromPublicKey(pub)
	if err != nil {
		return nil, fmt.Errorf("peer id from pubkey: %w", err)
	}
	boxKeys, err := GenerateBoxKey()
	if err != nil {
		return nil, fmt.Errorf("generate box key: %w", err)
	}
	token := make([]byte, 32)
	if _, err := rand.Read(token); err != nil {
		return nil, fmt.Errorf("generate auth token: %w", err)
	}
	return &OperatorKey{
		KeyPair:   KeyPair{PrivateKey: priv, PublicKey: pub},
		PeerID:    pid,
		BoxKeys:   boxKeys,
		AuthToken: token,
	}, nil
}

// GenerateImplantKey creates a fresh Ed25519 keypair and stores a reference to the
// operator's public key for signature verification.
func GenerateImplantKey(operatorPub crypto.PubKey) (*ImplantKey, error) {
	priv, pub, err := crypto.GenerateEd25519Key(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("generate implant key: %w", err)
	}
	pid, err := peer.IDFromPublicKey(pub)
	if err != nil {
		return nil, fmt.Errorf("peer id from pubkey: %w", err)
	}
	return &ImplantKey{
		KeyPair:        KeyPair{PrivateKey: priv, PublicKey: pub},
		PeerID:         pid,
		OperatorPubKey: operatorPub,
	}, nil
}

// LoadPrivateKey deserializes an Ed25519 private key from protobuf wire format.
func LoadPrivateKey(data []byte) (crypto.PrivKey, error) {
	return crypto.UnmarshalPrivateKey(data)
}

// MarshalPrivateKey serializes a private key to protobuf wire format.
func MarshalPrivateKey(priv crypto.PrivKey) ([]byte, error) {
	return crypto.MarshalPrivateKey(priv)
}

// PubKeyFromBytes deserializes a public key from protobuf wire format.
func PubKeyFromBytes(data []byte) (crypto.PubKey, error) {
	return crypto.UnmarshalPublicKey(data)
}

// Verify checks an Ed25519 signature against the given public key and data.
func Verify(pub crypto.PubKey, data []byte, sig []byte) (bool, error) {
	return pub.Verify(data, sig)
}

// BoxKeyPath returns the path to the operator's box private key file.
func BoxKeyPath(dir string) string {
	return filepath.Join(dir, "operator.boxkey")
}

// BoxPubKeyPath returns the path to the operator's box public key file.
func BoxPubKeyPath(dir string) string {
	return filepath.Join(dir, "operator.boxpub")
}

// AuthTokenPath returns the path to the operator's auth token file.
func AuthTokenPath(dir string) string {
	return filepath.Join(dir, "operator.authtoken")
}

// LoadOrGenerateOperatorKey tries to load an existing key pair from disk; if not found,
// generates a new one and persists it alongside the box keypair and auth token.
func LoadOrGenerateOperatorKey(path string) (*OperatorKey, error) {
	dir := filepath.Dir(path)
	if err := os.MkdirAll(dir, 0700); err != nil {
		return nil, fmt.Errorf("create key directory: %w", err)
	}

	data, err := os.ReadFile(path)
	if err == nil {
		priv, err := LoadPrivateKey(data)
		if err != nil {
			return nil, fmt.Errorf("load private key from %s: %w", path, err)
		}
		pub := priv.GetPublic()
		pid, err := peer.IDFromPublicKey(pub)
		if err != nil {
			return nil, fmt.Errorf("peer id from loaded key: %w", err)
		}

		boxPrivData, boxErr := os.ReadFile(BoxKeyPath(dir))
		var boxKeys *BoxKeyPair
		if boxErr == nil && len(boxPrivData) == 32 {
			boxKeys = LoadBoxKey(boxPrivData)
		}
		if boxKeys == nil {
			boxKeys, boxErr = GenerateBoxKey()
			if boxErr == nil {
				os.WriteFile(BoxKeyPath(dir), boxKeys.Marshal(), 0600)
				os.WriteFile(BoxPubKeyPath(dir), boxKeys.BoxPubKeyBytes(), 0644)
			}
		}

		tokenData, tokenErr := os.ReadFile(AuthTokenPath(dir))
		var authToken []byte
		if tokenErr == nil && len(tokenData) == 32 {
			authToken = tokenData
		} else {
			authToken = make([]byte, 32)
			if _, err := rand.Read(authToken); err != nil {
				return nil, fmt.Errorf("generate auth token: %w", err)
			}
			os.WriteFile(AuthTokenPath(dir), authToken, 0600)
		}

		return &OperatorKey{
			KeyPair:   KeyPair{PrivateKey: priv, PublicKey: pub},
			PeerID:    pid,
			BoxKeys:   boxKeys,
			AuthToken: authToken,
		}, nil
	}

	if !os.IsNotExist(err) {
		return nil, fmt.Errorf("read key file %s: %w", path, err)
	}

	key, err := GenerateOperatorKey()
	if err != nil {
		return nil, err
	}

	marshaled, err := MarshalPrivateKey(key.PrivateKey)
	if err != nil {
		return nil, fmt.Errorf("marshal key: %w", err)
	}

	if err := os.WriteFile(path, marshaled, 0600); err != nil {
		return nil, fmt.Errorf("write key to %s: %w", path, err)
	}
	if err := os.WriteFile(BoxKeyPath(dir), key.BoxKeys.Marshal(), 0600); err != nil {
		return nil, fmt.Errorf("write box key: %w", err)
	}
	if err := os.WriteFile(BoxPubKeyPath(dir), key.BoxKeys.BoxPubKeyBytes(), 0644); err != nil {
		return nil, fmt.Errorf("write box pubkey: %w", err)
	}
	if err := os.WriteFile(AuthTokenPath(dir), key.AuthToken, 0600); err != nil {
		return nil, fmt.Errorf("write auth token: %w", err)
	}

	return key, nil
}

// BoxPubKeyBytes returns the box public key as 32 raw bytes.
func (kp *BoxKeyPair) BoxPubKeyBytes() []byte {
	return kp.PublicKey[:]
}

// EncryptMessage encrypts plaintext using NaCl box anonymous sealed boxes.
func EncryptMessage(plaintext []byte, boxPub *[32]byte) ([]byte, error) {
	if len(plaintext) == 0 {
		return plaintext, nil
	}
	sealed, err := box.SealAnonymous(nil, plaintext, boxPub, rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("seal anonymous: %w", err)
	}
	return sealed, nil
}

// DecryptMessage decrypts a NaCl box anonymous sealed box using the provided keypair.
func DecryptMessage(ciphertext []byte, boxKeys *BoxKeyPair) ([]byte, error) {
	if len(ciphertext) == 0 {
		return ciphertext, nil
	}
	opened, ok := box.OpenAnonymous(nil, ciphertext, boxKeys.PublicKey, boxKeys.PrivateKey)
	if !ok {
		return nil, fmt.Errorf("open sealed box failed")
	}
	return opened, nil
}