// 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"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"time"

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

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

// b64 is the base64url encoder we use everywhere (no padding, URL-safe).
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 base64url encoding of the public key.
// This is the peer's stable identifier.
func (id *Identity) PeerID() proto.PeerID {
	return proto.PeerID(b64.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: b64.EncodeToString(id.PublicKey),
		CreatedAt: time.Now(),
	}
}

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

// Verify checks an Ed25519 signature against a base64url public key.
func Verify(publicKeyB64 string, data []byte, sigB64 string) error {
	pubBytes, err := b64.DecodeString(publicKeyB64)
	if err != nil {
		return fmt.Errorf("decoding public key: %w", err)
	}
	sigBytes, err := b64.DecodeString(sigB64)
	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
}

// ── 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[:])
}

// 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
}