package secp256k1

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

// Signature represents a Schnorr signature (r, s)
// r is the x-coordinate of R (32 bytes)
// s is the scalar response (32 bytes)
type Signature struct {
	R *big.Int // x-coordinate of the nonce point
	S *big.Int // the response scalar
}

// TaggedHash computes SHA256(SHA256(tag) || SHA256(tag) || msg)
// This is BIP-340's domain separation technique
func TaggedHash(tag string, data ...[]byte) []byte {
	tagHash := sha256.Sum256([]byte(tag))
	h := sha256.New()
	h.Write(tagHash[:])
	h.Write(tagHash[:])
	for _, d := range data {
		h.Write(d)
	}
	return h.Sum(nil)
}

// liftX recovers a point from just its x-coordinate
// Returns the point with even y (BIP-340 convention)
func liftX(x *big.Int) (*Point, error) {
	// Check x is in valid range
	if x.Sign() < 0 || x.Cmp(P) >= 0 {
		return nil, fmt.Errorf("x out of range")
	}

	// Compute y² = x³ + 7
	xFe := NewFieldElement(x)
	ySquared := xFe.Square().Mul(xFe).Add(curveB)

	// Compute y = sqrt(y²) mod p
	// For secp256k1, sqrt(a) = a^((p+1)/4) mod p
	exp := new(big.Int).Add(P, big.NewInt(1))
	exp.Div(exp, big.NewInt(4))
	y := new(big.Int).Exp(ySquared.value, exp, P)

	// Verify it's actually a square root
	ySquaredCheck := new(big.Int).Mul(y, y)
	ySquaredCheck.Mod(ySquaredCheck, P)
	if ySquaredCheck.Cmp(ySquared.value) != 0 {
		return nil, fmt.Errorf("x is not on the curve")
	}

	// BIP-340: use the even y
	if y.Bit(0) == 1 {
		y.Sub(P, y)
	}

	return &Point{
		x:        NewFieldElement(x),
		y:        NewFieldElement(y),
		infinity: false,
	}, nil
}

// hasEvenY returns true if the point's y-coordinate is even
func hasEvenY(p *Point) bool {
	if p.infinity {
		return false
	}
	return p.y.value.Bit(0) == 0
}

// xOnlyBytes returns the 32-byte x-coordinate of a public key
func (pub *PublicKey) XOnlyBytes() []byte {
	result := make([]byte, 32)
	xBytes := pub.Point.x.value.Bytes()
	copy(result[32-len(xBytes):], xBytes)
	return result
}

// Sign creates a Schnorr signature for a message
// Follows BIP-340 specification
// aux is optional auxiliary randomness (32 bytes); nil uses zeros
func Sign(priv *PrivateKey, message []byte, aux ...[]byte) (*Signature, error) {
	// Get the public key point
	P := G.ScalarMul(priv.D)

	// BIP-340: if P.y is odd, negate the private key
	d := new(big.Int).Set(priv.D)
	if !hasEvenY(P) {
		d.Sub(N, d)
		P = P.Negate()
	}

	// Serialize public key x-coordinate (32 bytes)
	pBytes := make([]byte, 32)
	pxBytes := P.x.value.Bytes()
	copy(pBytes[32-len(pxBytes):], pxBytes)

	// BIP-340 nonce generation:
	// t = d XOR tagged_hash("BIP0340/aux", aux)
	// k = tagged_hash("BIP0340/nonce", t || P || m)
	// For deterministic signing, use aux = 32 zero bytes
	dBytes := make([]byte, 32)
	dBytesRaw := d.Bytes()
	copy(dBytes[32-len(dBytesRaw):], dBytesRaw)

	// Use provided aux or default to 32 zero bytes
	var auxBytes []byte
	if len(aux) > 0 && len(aux[0]) == 32 {
		auxBytes = aux[0]
	} else {
		auxBytes = make([]byte, 32)
	}
	auxHash := TaggedHash("BIP0340/aux", auxBytes)

	t := make([]byte, 32)
	for i := 0; i < 32; i++ {
		t[i] = dBytes[i] ^ auxHash[i]
	}

	kHash := TaggedHash("BIP0340/nonce", t, pBytes, message)
	k := new(big.Int).SetBytes(kHash)
	k.Mod(k, N)

	// k cannot be zero (extremely unlikely)
	if k.Sign() == 0 {
		return nil, fmt.Errorf("nonce is zero")
	}

	// R = k * G
	R := G.ScalarMul(k)

	// BIP-340: if R.y is odd, negate k
	if !hasEvenY(R) {
		k.Sub(N, k)
		R = R.Negate()
	}

	// Serialize R.x (32 bytes)
	rBytes := make([]byte, 32)
	rxBytes := R.x.value.Bytes()
	copy(rBytes[32-len(rxBytes):], rxBytes)

	// Compute challenge e = hash(R.x || P.x || m)
	eHash := TaggedHash("BIP0340/challenge", rBytes, pBytes, message)
	e := new(big.Int).SetBytes(eHash)
	e.Mod(e, N)

	// Compute s = k + e * d (mod N)
	s := new(big.Int).Mul(e, d)
	s.Add(s, k)
	s.Mod(s, N)

	return &Signature{
		R: R.x.value,
		S: s,
	}, nil
}

// Verify checks if a Schnorr signature is valid
// Follows BIP-340 specification
func Verify(pub *PublicKey, message []byte, sig *Signature) bool {
	// Check signature values are in range
	if sig.R.Sign() < 0 || sig.R.Cmp(P) >= 0 {
		return false
	}
	if sig.S.Sign() < 0 || sig.S.Cmp(N) >= 0 {
		return false
	}

	// Lift R from x-coordinate
	R, err := liftX(sig.R)
	if err != nil {
		return false
	}

	// Get public key with even y
	P := pub.Point
	if !hasEvenY(P) {
		P = P.Negate()
	}

	// Serialize for hashing
	rBytes := make([]byte, 32)
	rxBytes := sig.R.Bytes()
	copy(rBytes[32-len(rxBytes):], rxBytes)

	pBytes := make([]byte, 32)
	pxBytes := P.x.value.Bytes()
	copy(pBytes[32-len(pxBytes):], pxBytes)

	// Compute challenge e = hash(R.x || P.x || m)
	eHash := TaggedHash("BIP0340/challenge", rBytes, pBytes, message)
	e := new(big.Int).SetBytes(eHash)
	e.Mod(e, N)

	// Verify: s*G == R + e*P
	sG := G.ScalarMul(sig.S)
	eP := P.ScalarMul(e)
	expected := R.Add(eP)

	return sG.Equal(expected)
}

// Bytes returns the signature as 64 bytes (r || s)
func (sig *Signature) Bytes() []byte {
	result := make([]byte, 64)

	rBytes := sig.R.Bytes()
	sBytes := sig.S.Bytes()

	copy(result[32-len(rBytes):32], rBytes)
	copy(result[64-len(sBytes):64], sBytes)

	return result
}

// SignatureFromBytes parses a 64-byte signature
func SignatureFromBytes(b []byte) (*Signature, error) {
	if len(b) != 64 {
		return nil, fmt.Errorf("signature must be 64 bytes")
	}

	r := new(big.Int).SetBytes(b[:32])
	s := new(big.Int).SetBytes(b[32:])

	return &Signature{R: r, S: s}, nil
}

// Hex returns the signature as a 128-character hex string
func (sig *Signature) Hex() string {
	return fmt.Sprintf("%x", sig.Bytes())
}