// Package crypto handles all cryptographic operations for waste-go.
//
// Identity:  Ed25519 keypair  — who you are, persistent across restarts
// Session:   X25519 ECDH     — per-connection shared secret
// Symmetric: ChaCha20-Poly1305 — encrypts every message after the handshake
package crypto

import (
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"time"

	"filippo.io/edwards25519"
	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/curve25519"
	"golang.org/x/crypto/hkdf"
	"golang.org/x/crypto/nacl/box"

	"github.com/waste-go/internal/proto"
)

// b64 is used only for storing the private key on disk and for symmetric nonces/ciphertexts.
var b64 = base64.RawURLEncoding

// ── Identity ──────────────────────────────────────────────────────────────────

// Identity holds our Ed25519 keypair and alias.
// Load it once at startup with LoadOrCreate; then pass it around.
type Identity struct {
	privateKey ed25519.PrivateKey
	PublicKey  ed25519.PublicKey
	Alias      string
}

// identityFile is what we store on disk.
type identityFile struct {
	PrivateKeyB64 string `json:"private_key"`
	Alias         string `json:"alias"`
}

// LoadOrCreate loads the identity from dataDir/identity.json,
// or generates a fresh keypair if none exists yet.
func LoadOrCreate(dataDir, alias string) (*Identity, error) {
	if err := os.MkdirAll(dataDir, 0700); err != nil {
		return nil, fmt.Errorf("creating data dir: %w", err)
	}

	path := filepath.Join(dataDir, "identity.json")

	data, err := os.ReadFile(path)
	if err == nil {
		// File exists — load it
		var f identityFile
		if err := json.Unmarshal(data, &f); err != nil {
			return nil, fmt.Errorf("parsing identity file: %w", err)
		}
		privBytes, err := b64.DecodeString(f.PrivateKeyB64)
		if err != nil {
			return nil, fmt.Errorf("decoding private key: %w", err)
		}
		priv := ed25519.PrivateKey(privBytes)
		return &Identity{
			privateKey: priv,
			PublicKey:  priv.Public().(ed25519.PublicKey),
			Alias:      f.Alias,
		}, nil
	}

	if !errors.Is(err, os.ErrNotExist) {
		return nil, fmt.Errorf("reading identity file: %w", err)
	}

	// Generate new keypair
	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("generating keypair: %w", err)
	}

	f := identityFile{
		PrivateKeyB64: b64.EncodeToString(priv),
		Alias:         alias,
	}
	raw, _ := json.MarshalIndent(f, "", "  ")
	if err := os.WriteFile(path, raw, 0600); err != nil {
		return nil, fmt.Errorf("saving identity: %w", err)
	}

	fmt.Printf("Generated new identity, saved to %s\n", path)
	return &Identity{privateKey: priv, PublicKey: pub, Alias: alias}, nil
}

// PeerID returns the lowercase hex encoding of the 32-byte Ed25519 public key (YAW/2 §2).
func (id *Identity) PeerID() proto.PeerID {
	return proto.PeerID(hex.EncodeToString(id.PublicKey))
}

// PeerInfo builds the PeerInfo struct for the handshake.
func (id *Identity) PeerInfo() proto.PeerInfo {
	return proto.PeerInfo{
		ID:        id.PeerID(),
		Alias:     id.Alias,
		PublicKey: hex.EncodeToString(id.PublicKey),
		CreatedAt: time.Now(),
	}
}

// Sign signs data with our Ed25519 private key. Returns hex-encoded signature.
func (id *Identity) Sign(data []byte) string {
	sig := ed25519.Sign(id.privateKey, data)
	return hex.EncodeToString(sig)
}

// Verify checks an Ed25519 signature against a hex-encoded public key.
func Verify(publicKeyHex string, data []byte, sigHex string) error {
	pubBytes, err := hex.DecodeString(publicKeyHex)
	if err != nil {
		return fmt.Errorf("decoding public key: %w", err)
	}
	sigBytes, err := hex.DecodeString(sigHex)
	if err != nil {
		return fmt.Errorf("decoding signature: %w", err)
	}
	if !ed25519.Verify(ed25519.PublicKey(pubBytes), data, sigBytes) {
		return errors.New("signature verification failed")
	}
	return nil
}

// CurvePublicKey returns the X25519 public key derived from this Ed25519 identity (YAW/2 §3).
func (id *Identity) CurvePublicKey() *[32]byte {
	edPoint, _ := new(edwards25519.Point).SetBytes(id.PublicKey)
	mont := edPoint.BytesMontgomery()
	var out [32]byte
	copy(out[:], mont)
	return &out
}

// CurvePrivateKey returns the X25519 private key derived from this Ed25519 identity (YAW/2 §3).
// Uses the clamped SHA-512 first half of the Ed25519 seed, matching libsodium's conversion.
func (id *Identity) CurvePrivateKey() *[32]byte {
	h := sha512.Sum512(id.privateKey[:32])
	h[0] &= 248
	h[31] &= 127
	h[31] |= 64
	var out [32]byte
	copy(out[:], h[:32])
	return &out
}

// SignalingBox seals a plaintext signaling payload for a recipient using nacl/box (YAW/2 §5).
// Returns base64(nonce(24) || ciphertext).
func SignalingBox(plaintext []byte, recipientPub, senderPriv *[32]byte) string {
	var nonce [24]byte
	rand.Read(nonce[:]) //nolint:errcheck
	ct := box.Seal(nonce[:], plaintext, &nonce, recipientPub, senderPriv)
	return base64.StdEncoding.EncodeToString(ct)
}

// SignalingOpen opens a nacl/box sealed by SignalingBox.
func SignalingOpen(b64box string, senderPub, recipientPriv *[32]byte) ([]byte, error) {
	raw, err := base64.StdEncoding.DecodeString(b64box)
	if err != nil || len(raw) < 24 {
		return nil, errors.New("invalid box")
	}
	var nonce [24]byte
	copy(nonce[:], raw[:24])
	out, ok := box.Open(nil, raw[24:], &nonce, senderPub, recipientPriv)
	if !ok {
		return nil, errors.New("box open failed")
	}
	return out, nil
}

// DeriveForNetwork returns a new in-memory Identity derived from the master key
// and the given network hash (hex). Same master + same network hash always
// produces the same keypair, so the peer ID is stable across restarts.
// Different networks produce different peer IDs, preventing cross-network correlation.
func DeriveForNetwork(master *Identity, networkHash string) (*Identity, error) {
	r := hkdf.New(sha256.New, master.privateKey[:32], []byte(networkHash), []byte("yaw2-net-identity"))
	var seed [32]byte
	if _, err := io.ReadFull(r, seed[:]); err != nil {
		return nil, fmt.Errorf("deriving network identity: %w", err)
	}
	priv := ed25519.NewKeyFromSeed(seed[:])
	return &Identity{
		privateKey: priv,
		PublicKey:  priv.Public().(ed25519.PublicKey),
		Alias:      master.Alias,
	}, nil
}

// ── X25519 ECDH ───────────────────────────────────────────────────────────────

// EphemeralKey is an X25519 keypair used for a single session.
type EphemeralKey struct {
	private [32]byte
	public  [32]byte
}

// GenerateEphemeral creates a fresh X25519 keypair.
func GenerateEphemeral() (*EphemeralKey, error) {
	ek := &EphemeralKey{}
	if _, err := rand.Read(ek.private[:]); err != nil {
		return nil, fmt.Errorf("generating ephemeral key: %w", err)
	}
	// Clamp per RFC 7748
	ek.private[0] &= 248
	ek.private[31] &= 127
	ek.private[31] |= 64

	pub, err := curve25519.X25519(ek.private[:], curve25519.Basepoint)
	if err != nil {
		return nil, fmt.Errorf("computing public key: %w", err)
	}
	copy(ek.public[:], pub)
	return ek, nil
}

// PublicKeyB64 returns the base64url public key to send in the handshake.
func (ek *EphemeralKey) PublicKeyB64() string {
	return b64.EncodeToString(ek.public[:])
}

// PublicRaw returns a pointer to the raw 32-byte X25519 public key.
// The returned pointer is valid for the lifetime of the EphemeralKey.
func (ek *EphemeralKey) PublicRaw() *[32]byte { return &ek.public }

// PrivateRaw returns a pointer to the raw 32-byte X25519 private key.
// Use only when you need to pass it directly to SignalingBox.
func (ek *EphemeralKey) PrivateRaw() *[32]byte { return &ek.private }

// Wipe zeroes the private key. Call when the session ends.
func (ek *EphemeralKey) Wipe() {
	for i := range ek.private {
		ek.private[i] = 0
	}
}

// SharedSecret performs ECDH with the other party's public key.
// Returns a 32-byte shared secret suitable for use as an AEAD key.
func (ek *EphemeralKey) SharedSecret(theirPublicB64 string) ([32]byte, error) {
	var result [32]byte
	theirBytes, err := b64.DecodeString(theirPublicB64)
	if err != nil {
		return result, fmt.Errorf("decoding their public key: %w", err)
	}
	shared, err := curve25519.X25519(ek.private[:], theirBytes)
	if err != nil {
		return result, fmt.Errorf("ECDH: %w", err)
	}
	// Hash the raw shared secret (simple KDF — use HKDF in production)
	result = sha256.Sum256(shared)
	return result, nil
}

// ── ChaCha20-Poly1305 AEAD ────────────────────────────────────────────────────

// Session holds the symmetric key for an established peer session.
type Session struct {
	key [32]byte
}

// NewSession wraps a 32-byte shared secret as a usable session.
func NewSession(key [32]byte) *Session {
	return &Session{key: key}
}

// Encrypt encrypts plaintext. Returns (nonce_b64, ciphertext_b64).
func (s *Session) Encrypt(plaintext []byte) (string, string, error) {
	aead, err := chacha20poly1305.New(s.key[:])
	if err != nil {
		return "", "", fmt.Errorf("creating AEAD: %w", err)
	}

	nonce := make([]byte, aead.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return "", "", fmt.Errorf("generating nonce: %w", err)
	}

	ciphertext := aead.Seal(nil, nonce, plaintext, nil)
	return b64.EncodeToString(nonce), b64.EncodeToString(ciphertext), nil
}

// Decrypt decrypts a nonce+ciphertext pair. Returns plaintext.
func (s *Session) Decrypt(nonceB64, ciphertextB64 string) ([]byte, error) {
	aead, err := chacha20poly1305.New(s.key[:])
	if err != nil {
		return nil, fmt.Errorf("creating AEAD: %w", err)
	}

	nonce, err := b64.DecodeString(nonceB64)
	if err != nil {
		return nil, fmt.Errorf("decoding nonce: %w", err)
	}
	ciphertext, err := b64.DecodeString(ciphertextB64)
	if err != nil {
		return nil, fmt.Errorf("decoding ciphertext: %w", err)
	}

	plaintext, err := aead.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, errors.New("decryption failed — bad key or tampered message")
	}
	return plaintext, nil
}