package pow

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/allegro/bigcache/v3"
	"github.com/google/uuid"
	"github.com/rs/zerolog/log"

	"pluja.dev/kycnot.me/utils"
)

type PowChallenger struct {
	IncreaseEveryChallenges int
	Difficulty              int
	Cache                   *bigcache.BigCache
}

func (p *PowChallenger) Init() {
	// Parse the environment variable once, at initialization
	log.Printf("Initializing PoW challenger.")
	if iec := utils.Getenv("POW_INCREASE_EVERY_CHALLENGES", "20"); iec != "" {
		if parsed, err := strconv.Atoi(iec); err == nil {
			p.IncreaseEveryChallenges = parsed
		}
	}
	p.IncreaseEveryChallenges = 20

	var err error
	p.Difficulty, err = strconv.Atoi(utils.Getenv("POW_DIFFICULTY", "4"))
	if err != nil {
		log.Error().Err(err).Msg("Could not get difficulty from environment")
		p.Difficulty = 4
	}

	// Init cache
	cache, _ := bigcache.New(context.Background(), bigcache.Config{
		Shards:             1024,
		LifeWindow:         10 * time.Minute,
		CleanWindow:        2 * time.Minute,
		MaxEntriesInWindow: 1000 * 10 * 60,
		MaxEntrySize:       500,
		Verbose:            false,
		HardMaxCacheSize:   0,
	})
	p.Cache = cache
}

// GenerateChallenge generates a new challenge for the client.
func (p *PowChallenger) PowGenerateChallenge(length int) (challenge, id string, difficulty int, err error) {
	bytes := make([]byte, length/2) // since hex encoding doubles the length
	_, err = rand.Read(bytes)
	if err != nil {
		return "", "", 0, err
	}
	id = uuid.New().String()
	challenge = hex.EncodeToString(bytes)

	// Increase the difficulty by 1 every N challenges in a 10 minute period
	count := p.Cache.Len()
	difficulty = p.Difficulty
	if count%p.IncreaseEveryChallenges == 0 && count != 0 {
		difficulty++
	}

	log.Printf("Generated challenge %v with difficulty %v", id, difficulty)

	// Store the challenge in a temporary cache
	err = p.Cache.Set(fmt.Sprintf("pow-%v-c", id), bytes)
	if err != nil {
		log.Error().Err(err).Msg("Could not set challenge in cache")
		return "", "", 0, err
	}

	// Store the challenge in a temporary cache
	difficultyBytes := make([]byte, 1)
	difficultyBytes[0] = byte(difficulty)
	err = p.Cache.Set(fmt.Sprintf("pow-%v-d", id), difficultyBytes)
	if err != nil {
		log.Error().Err(err).Msg("Could not set difficulty in cache")
		return "", "", 0, err
	}

	return
}

// VerifyProof verifies the client's proof of work.
func (p *PowChallenger) PowVerifyProof(id, nonce string) bool {
	// Get the challenge from the cache
	challengeBytes, err := p.Cache.Get(fmt.Sprintf("pow-%v-c", id))
	if err != nil {
		log.Error().Err(err).Msg("Could not get challenge from cache")
		return false
	}
	challenge := hex.EncodeToString(challengeBytes)

	// Get the difficulty from the cache
	difficultyBytes, err := p.Cache.Get(fmt.Sprintf("pow-%v-d", id))
	if err != nil {
		log.Error().Err(err).Msg("Could not get difficulty from cache")
		return false
	}
	difficulty := int(difficultyBytes[0])

	candidate := fmt.Sprintf("%s:%s", challenge, nonce)
	hash := sha256.Sum256([]byte(candidate))
	hashHex := hex.EncodeToString(hash[:])

	return strings.HasPrefix(hashHex, strings.Repeat("0", difficulty))
}
