package snapshot

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"sync"
	"sync/atomic"

	proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy"
)

var (
	ErrWrongTarget      = errors.New("snapshot targets another cluster or worker")
	ErrResyncRequired   = errors.New("snapshot sequence requires a full resync")
	ErrChecksumMismatch = errors.New("snapshot checksum mismatch")
)

type Envelope struct {
	ClusterID string
	WorkerID  string
	Epoch     uint64
	Version   uint64
	Full      bool
	Checksum  string
	Proxies   []proxyDomain.Proxy
}

type Entry struct {
	Proxy   proxyDomain.Proxy
	Runtime *proxyDomain.Capacity
}

type View struct {
	ClusterID string
	WorkerID  string
	Epoch     uint64
	Version   uint64
	Checksum  string
	Entries   []Entry
}

type Store struct {
	clusterID string
	workerID  string
	current   atomic.Pointer[View]

	mu       sync.Mutex
	runtimes map[string]*proxyDomain.Capacity
}

func NewStore(clusterID, workerID string) *Store {
	return &Store{
		clusterID: clusterID,
		workerID:  workerID,
		runtimes:  make(map[string]*proxyDomain.Capacity),
	}
}

func (s *Store) Current() *View {
	if s == nil {
		return nil
	}
	return s.current.Load()
}

func (s *Store) Apply(envelope Envelope) error {
	if s == nil {
		return fmt.Errorf("apply snapshot: nil store")
	}
	if envelope.ClusterID != s.clusterID || envelope.WorkerID != s.workerID {
		return ErrWrongTarget
	}
	if !envelope.Full || envelope.Epoch == 0 || envelope.Version == 0 {
		return ErrResyncRequired
	}
	if envelope.Checksum != Checksum(envelope.Proxies) {
		return ErrChecksumMismatch
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	current := s.current.Load()
	if current != nil {
		switch {
		case envelope.Epoch < current.Epoch:
			return ErrResyncRequired
		case envelope.Epoch == current.Epoch && envelope.Version != current.Version+1:
			return ErrResyncRequired
		case envelope.Epoch > current.Epoch && envelope.Version != 1:
			return ErrResyncRequired
		}
	}

	proxies := cloneAndSort(envelope.Proxies)
	entries := make([]Entry, 0, len(proxies))
	for _, descriptor := range proxies {
		runtime := s.runtimes[descriptor.ID]
		if runtime == nil {
			runtime = proxyDomain.NewCapacity(descriptor.MaxConcurrency)
			s.runtimes[descriptor.ID] = runtime
		} else {
			runtime.SetMax(descriptor.MaxConcurrency)
		}
		entries = append(entries, Entry{Proxy: descriptor, Runtime: runtime})
	}

	next := &View{
		ClusterID: envelope.ClusterID,
		WorkerID:  envelope.WorkerID,
		Epoch:     envelope.Epoch,
		Version:   envelope.Version,
		Checksum:  envelope.Checksum,
		Entries:   entries,
	}
	s.current.Store(next)
	return nil
}

func Checksum(proxies []proxyDomain.Proxy) string {
	canonical := cloneAndSort(proxies)
	encoded, err := json.Marshal(canonical)
	if err != nil {
		panic(fmt.Sprintf("encode snapshot checksum: %v", err))
	}
	digest := sha256.Sum256(encoded)
	return hex.EncodeToString(digest[:])
}

func cloneAndSort(source []proxyDomain.Proxy) []proxyDomain.Proxy {
	cloned := make([]proxyDomain.Proxy, len(source))
	for index, descriptor := range source {
		cloned[index] = descriptor
		if descriptor.Tags != nil {
			cloned[index].Tags = make(map[string]string, len(descriptor.Tags))
			for key, value := range descriptor.Tags {
				cloned[index].Tags[key] = value
			}
		}
	}
	sort.Slice(cloned, func(i, j int) bool {
		return cloned[i].ID < cloned[j].ID
	})
	return cloned
}