// Package wallet manages Ed25519 keypairs with human-readable addresses,
// encrypted persistence, and two usage profiles:
//
//   - NodeWallet  — bound to a running validator node; earns block/relay rewards.
//   - UserWallet  — holds tokens for regular users; spends on transactions.
//
// Address format: "DC" + lowercase hex(sha256(pubkey)[0:12])
// Example: DC3a7f2b1c9d0e5f6a7b8c9dab1f
// (2 + 24 = 26 characters, collision-resistant for this use case)
package wallet

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"

	"golang.org/x/crypto/argon2"

	"go-blockchain/identity"
)

// WalletType indicates how the wallet is used.
type WalletType string

const (
	NodeWallet WalletType = "NODE" // validator node; earns rewards
	UserWallet WalletType = "USER" // end-user; spends tokens
)

// Wallet holds an Ed25519 identity plus metadata.
type Wallet struct {
	Type    WalletType
	ID      *identity.Identity
	Label   string // human-readable nickname
	Address string // "DC..." address derived from pub key
}

// New creates a fresh wallet of the given type and label.
func New(wtype WalletType, label string) (*Wallet, error) {
	id, err := identity.Generate()
	if err != nil {
		return nil, err
	}
	return &Wallet{
		Type:    wtype,
		ID:      id,
		Label:   label,
		Address: PubKeyToAddress(id.PubKeyHex()),
	}, nil
}

// FromIdentity wraps an existing identity in a Wallet.
func FromIdentity(id *identity.Identity, wtype WalletType, label string) *Wallet {
	return &Wallet{
		Type:    wtype,
		ID:      id,
		Label:   label,
		Address: PubKeyToAddress(id.PubKeyHex()),
	}
}

// PubKeyToAddress derives the "DC…" address from a hex-encoded Ed25519 public key.
func PubKeyToAddress(pubKeyHex string) string {
	raw, err := hex.DecodeString(pubKeyHex)
	if err != nil {
		return "DCinvalid"
	}
	h := sha256.Sum256(raw)
	return "DC" + hex.EncodeToString(h[:12])
}

// AddressToPubKeyPrefix returns the first 12 bytes of the address payload (for display).
func AddressToPubKeyPrefix(addr string) string {
	if len(addr) < 2 {
		return addr
	}
	return addr[2:] // strip "DC" prefix
}

// --- persistence ---

// walletFile is the JSON structure saved to disk.
// The private key is stored AES-256-GCM encrypted with Argon2id key derivation.
type walletFile struct {
	Version int    `json:"version"`
	Type    string `json:"type"`
	Label   string `json:"label"`
	Address string `json:"address"`
	PubKey  string `json:"pub_key"`

	// Encryption envelope
	KDFSalt  string `json:"kdf_salt"`  // hex-encoded 16-byte Argon2id salt
	Nonce    string `json:"nonce"`     // hex-encoded 12-byte AES-GCM nonce
	EncPriv  string `json:"enc_priv"`  // hex-encoded AES-256-GCM ciphertext of priv key
}

// Save encrypts the wallet and writes it to path.
// passphrase may be "" (unencrypted, not recommended for production).
func (w *Wallet) Save(path, passphrase string) error {
	privKeyBytes := []byte(hex.EncodeToString(w.ID.PrivKey))

	salt := make([]byte, 16)
	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
		return fmt.Errorf("generate salt: %w", err)
	}
	nonce := make([]byte, 12)
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return fmt.Errorf("generate nonce: %w", err)
	}

	key := deriveKey(passphrase, salt)
	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return err
	}
	ct := gcm.Seal(nil, nonce, privKeyBytes, nil)

	wf := walletFile{
		Version: 1,
		Type:    string(w.Type),
		Label:   w.Label,
		Address: w.Address,
		PubKey:  w.ID.PubKeyHex(),
		KDFSalt: hex.EncodeToString(salt),
		Nonce:   hex.EncodeToString(nonce),
		EncPriv: hex.EncodeToString(ct),
	}
	data, err := json.MarshalIndent(wf, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(path, data, 0600)
}

// Load decrypts a wallet file.
func Load(path, passphrase string) (*Wallet, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read wallet: %w", err)
	}
	var wf walletFile
	if err := json.Unmarshal(data, &wf); err != nil {
		return nil, fmt.Errorf("parse wallet: %w", err)
	}

	salt, err := hex.DecodeString(wf.KDFSalt)
	if err != nil {
		return nil, fmt.Errorf("decode salt: %w", err)
	}
	nonce, err := hex.DecodeString(wf.Nonce)
	if err != nil {
		return nil, fmt.Errorf("decode nonce: %w", err)
	}
	ct, err := hex.DecodeString(wf.EncPriv)
	if err != nil {
		return nil, fmt.Errorf("decode ciphertext: %w", err)
	}

	key := deriveKey(passphrase, salt)
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	privKeyHex, err := gcm.Open(nil, nonce, ct, nil)
	if err != nil {
		return nil, errors.New("decryption failed: wrong passphrase?")
	}

	id, err := identity.FromHex(wf.PubKey, string(privKeyHex))
	if err != nil {
		return nil, fmt.Errorf("load identity: %w", err)
	}
	return &Wallet{
		Type:    WalletType(wf.Type),
		ID:      id,
		Label:   wf.Label,
		Address: wf.Address,
	}, nil
}

// deriveKey derives a 32-byte AES key from passphrase + salt using Argon2id.
func deriveKey(passphrase string, salt []byte) []byte {
	// Argon2id parameters: time=1, memory=64MB, threads=4
	return argon2.IDKey([]byte(passphrase), salt, 1, 64*1024, 4, 32)
}

// --- display helpers ---

// Info returns a map suitable for JSON marshalling.
func (w *Wallet) Info() map[string]interface{} {
	return map[string]interface{}{
		"type":    string(w.Type),
		"label":   w.Label,
		"address": w.Address,
		"pub_key": w.ID.PubKeyHex(),
	}
}

// Short returns "label (DCxxxx…xxxx)".
func (w *Wallet) Short() string {
	addr := w.Address
	if len(addr) > 10 {
		addr = addr[:6] + "…" + addr[len(addr)-4:]
	}
	return fmt.Sprintf("%s (%s)", w.Label, addr)
}