package secp256k1

import (
	"crypto/rand"
	"fmt"
	"math/big"
)

// PrivateKey is a scalar (1 to N-1) used for signing
type PrivateKey struct {
	D *big.Int // the secret scalar
}

// PublicKey is a point on the curve (D * G)
type PublicKey struct {
	Point *Point
}

// GeneratePrivateKey creates a random private key
func GeneratePrivateKey() (*PrivateKey, error) {
	// Generate random bytes
	bytes := make([]byte, 32)
	_, err := rand.Read(bytes)
	if err != nil {
		return nil, fmt.Errorf("failed to generate random bytes: %w", err)
	}

	// Convert to big.Int and reduce mod N
	d := new(big.Int).SetBytes(bytes)
	d.Mod(d, N)

	// Ensure it's not zero (extremely unlikely but must check)
	if d.Sign() == 0 {
		d.SetInt64(1)
	}

	return &PrivateKey{D: d}, nil
}

// NewPrivateKeyFromBytes creates a private key from 32 bytes
func NewPrivateKeyFromBytes(b []byte) (*PrivateKey, error) {
	if len(b) != 32 {
		return nil, fmt.Errorf("private key must be 32 bytes, got %d", len(b))
	}

	d := new(big.Int).SetBytes(b)

	// Validate: must be in range [1, N-1]
	if d.Sign() == 0 {
		return nil, fmt.Errorf("private key cannot be zero")
	}
	if d.Cmp(N) >= 0 {
		return nil, fmt.Errorf("private key must be less than curve order N")
	}

	return &PrivateKey{D: d}, nil
}

// NewPrivateKeyFromHex creates a private key from a hex string
func NewPrivateKeyFromHex(hex string) (*PrivateKey, error) {
	d, ok := new(big.Int).SetString(hex, 16)
	if !ok {
		return nil, fmt.Errorf("invalid hex string")
	}

	// Validate range
	if d.Sign() == 0 {
		return nil, fmt.Errorf("private key cannot be zero")
	}
	if d.Cmp(N) >= 0 {
		return nil, fmt.Errorf("private key must be less than curve order N")
	}

	return &PrivateKey{D: d}, nil
}

// PublicKey derives the public key from the private key
// PublicKey = D * G
func (priv *PrivateKey) PublicKey() *PublicKey {
	point := G.ScalarMul(priv.D)
	return &PublicKey{Point: point}
}

// Bytes returns the private key as 32 bytes (big-endian, zero-padded)
func (priv *PrivateKey) Bytes() []byte {
	b := priv.D.Bytes()
	// Pad to 32 bytes
	if len(b) < 32 {
		padded := make([]byte, 32)
		copy(padded[32-len(b):], b)
		return padded
	}
	return b
}

// Hex returns the private key as a 64-character hex string
func (priv *PrivateKey) Hex() string {
	return fmt.Sprintf("%064x", priv.D)
}

// Bytes returns the public key in uncompressed format (65 bytes: 0x04 || x || y)
func (pub *PublicKey) Bytes() []byte {
	if pub.Point.IsInfinity() {
		return []byte{0x00} // shouldn't happen with valid keys
	}

	result := make([]byte, 65)
	result[0] = 0x04 // uncompressed prefix

	xBytes := pub.Point.x.value.Bytes()
	yBytes := pub.Point.y.value.Bytes()

	// Copy x (padded to 32 bytes)
	copy(result[1+(32-len(xBytes)):33], xBytes)
	// Copy y (padded to 32 bytes)
	copy(result[33+(32-len(yBytes)):65], yBytes)

	return result
}

// BytesCompressed returns the public key in compressed format (33 bytes: prefix || x)
// Prefix is 0x02 if y is even, 0x03 if y is odd
func (pub *PublicKey) BytesCompressed() []byte {
	if pub.Point.IsInfinity() {
		return []byte{0x00}
	}

	result := make([]byte, 33)

	// Prefix based on y parity
	if pub.Point.y.value.Bit(0) == 0 {
		result[0] = 0x02 // y is even
	} else {
		result[0] = 0x03 // y is odd
	}

	xBytes := pub.Point.x.value.Bytes()
	copy(result[1+(32-len(xBytes)):33], xBytes)

	return result
}

// Hex returns the public key as uncompressed hex (130 characters)
func (pub *PublicKey) Hex() string {
	return fmt.Sprintf("%x", pub.Bytes())
}

// HexCompressed returns the public key as compressed hex (66 characters)
func (pub *PublicKey) HexCompressed() string {
	return fmt.Sprintf("%x", pub.BytesCompressed())
}

// Equal checks if two public keys are the same
func (pub *PublicKey) Equal(other *PublicKey) bool {
	return pub.Point.Equal(other.Point)
}