package forge

import (
	"crypto/ed25519"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"os"
)

// KeyPair holds an ed25519 signing key.
type KeyPair struct {
	Private ed25519.PrivateKey
	Public  ed25519.PublicKey
}

// LoadOrCreateKey loads an ed25519 key from disk or generates a new one.
func LoadOrCreateKey(path string) (*KeyPair, error) {
	data, err := os.ReadFile(path)
	if err == nil {
		if len(data) == ed25519.PrivateKeySize {
			priv := ed25519.PrivateKey(data)
			return &KeyPair{Private: priv, Public: priv.Public().(ed25519.PublicKey)}, nil
		}
		if len(data) == ed25519.SeedSize {
			priv := ed25519.NewKeyFromSeed(data)
			return &KeyPair{Private: priv, Public: priv.Public().(ed25519.PublicKey)}, nil
		}
		return nil, fmt.Errorf("invalid key size %d", len(data))
	}
	if !os.IsNotExist(err) {
		return nil, fmt.Errorf("read key: %w", err)
	}

	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("generate key: %w", err)
	}

	if err := os.WriteFile(path, priv.Seed(), 0o600); err != nil {
		return nil, fmt.Errorf("write key: %w", err)
	}

	return &KeyPair{Private: priv, Public: pub}, nil
}

// Sign returns a base64-encoded ed25519 signature of data.
func (k *KeyPair) Sign(data []byte) string {
	sig := ed25519.Sign(k.Private, data)
	return base64.StdEncoding.EncodeToString(sig)
}

// Verify checks a base64 signature against data using the public key.
func Verify(pub ed25519.PublicKey, data []byte, signatureB64 string) bool {
	sig, err := base64.StdEncoding.DecodeString(signatureB64)
	if err != nil {
		return false
	}
	return ed25519.Verify(pub, data, sig)
}

// PublicKeyHex returns the hex-encoded public key for embedding in agents.
func (k *KeyPair) PublicKeyHex() string {
	return hex.EncodeToString(k.Public)
}