package pow

import (
	"crypto/hmac"
	"crypto/sha256"
	"errors"
	"fmt"
	"time"
)

// PoW package specific errors
var (
	ErrInvalidHMAC        = errors.New("challenge HMAC signature is invalid")
	ErrExpiredChallenge   = errors.New("challenge has expired")
	ErrInvalidSolution    = errors.New("proof of work solution is invalid")
	ErrInvalidDifficulty  = errors.New("difficulty level is invalid")
	ErrMalformedChallenge = errors.New("challenge format is malformed")
	ErrInvalidConfig      = errors.New("configuration is invalid")
)

// Challenge represents a Proof of Work challenge with HMAC authentication
type Challenge struct {
	Timestamp  int64  `json:"timestamp"`
	Difficulty int    `json:"difficulty"`
	Resource   string `json:"resource"`
	Random     []byte `json:"random"`
	HMAC       []byte `json:"hmac"`
}

// CanonicalBytes returns the canonical byte representation for HMAC signing
func (c *Challenge) CanonicalBytes() []byte {
	// Use consistent field ordering: timestamp:difficulty:resource:random
	buf := make([]byte, 0, 64)
	buf = fmt.Appendf(buf, "%d:%d:%s:%x",
		c.Timestamp,
		c.Difficulty,
		c.Resource,
		c.Random,
	)
	return buf
}

// BuildSolutionString creates the string used for PoW computation
func (c *Challenge) BuildSolutionString(nonce uint64) []byte {
	// Format: resource:timestamp:difficulty:random:nonce
	buf := make([]byte, 0, 64)
	buf = fmt.Appendf(buf, "%s:%d:%d:%x:%d",
		c.Resource,
		c.Timestamp,
		c.Difficulty,
		c.Random,
		nonce,
	)
	return buf
}

// VerifySolution verifies if a solution is valid for this challenge
func (c *Challenge) VerifySolution(nonce uint64) bool {
	// Build solution string
	solutionStr := c.BuildSolutionString(nonce)

	// Compute SHA-256 hash
	hash := sha256.Sum256(solutionStr)

	// Check if hash has required leading zero bits
	return hasLeadingZeroBits(hash[:], c.Difficulty)
}

// hasLeadingZeroBits checks if hash has the required number of leading zero bits
func hasLeadingZeroBits(hash []byte, difficulty int) bool {
	full := difficulty >> 3        // number of whole zero bytes
	rem  := uint(difficulty & 7)   // remaining leading zero bits

	for i := range full {
		if hash[i] != 0 {
			return false
		}
	}
	if rem == 0 {
		return true
	}
	mask := byte(0xFF) << (8 - rem) // e.g., rem=3 => 11100000
	return (hash[full] & mask) == 0
}

// IsExpired checks if challenge has exceeded TTL
func (c *Challenge) IsExpired(ttl time.Duration) bool {
	challengeTime := time.Unix(c.Timestamp, 0)
	return time.Since(challengeTime) > ttl
}

// Sign creates and sets HMAC signature for the challenge
func (c *Challenge) Sign(secret []byte) {
	// Create canonical bytes for signing (excluding HMAC field)
	canonical := c.CanonicalBytes()

	// Create HMAC
	h := hmac.New(sha256.New, secret)
	h.Write(canonical)
	c.HMAC = h.Sum(nil)
}

// VerifyHMAC verifies the HMAC signature of the challenge
func (c *Challenge) VerifyHMAC(secret []byte) error {
	if len(c.HMAC) == 0 {
		return ErrMalformedChallenge
	}

	// Create canonical bytes (excluding HMAC)
	canonical := c.CanonicalBytes()

	// Compute expected HMAC
	h := hmac.New(sha256.New, secret)
	h.Write(canonical)
	expectedSignature := h.Sum(nil)

	// Compare signatures
	if !hmac.Equal(c.HMAC, expectedSignature) {
		return ErrInvalidHMAC
	}

	return nil
}

// Solution represents a client's solution to a PoW challenge
type Solution struct {
	Challenge Challenge `json:"challenge"`
	Nonce     uint64    `json:"nonce"`
}

// Verify verifies if this solution is valid
func (s *Solution) Verify() bool {
	return s.Challenge.VerifySolution(s.Nonce)
}