package provider

import (
	"errors"
	"math"
	"sync"

	"proxy-pool/internal/domain/upstream"
)

var ErrInvalidStatsRecorder = errors.New("invalid Provider stats recorder")

type Stats struct {
	UpstreamID            string
	ConsecutiveEmptyFetch int64
	FetchErrorCount       int64
}

type StatsReader interface {
	ReadProviderStats([]string) []Stats
}

type StatsRetainer interface {
	RetainProviderStats([]string)
}

type StatsRecorder struct {
	mu      sync.Mutex
	maximum int
	byID    map[string]Stats
}

func NewStatsRecorder(maximum int) (*StatsRecorder, error) {
	if maximum <= 0 {
		return nil, ErrInvalidStatsRecorder
	}
	return &StatsRecorder{maximum: maximum, byID: make(map[string]Stats)}, nil
}

func (recorder *StatsRecorder) Record(result Result) {
	if recorder == nil || result.UpstreamID == "" {
		return
	}
	recorder.mu.Lock()
	defer recorder.mu.Unlock()
	stats, exists := recorder.byID[result.UpstreamID]
	if !exists {
		if len(recorder.byID) >= recorder.maximum {
			return
		}
		stats.UpstreamID = result.UpstreamID
	}
	switch result.Class {
	case upstream.FetchEmpty:
		if stats.ConsecutiveEmptyFetch < math.MaxInt64 {
			stats.ConsecutiveEmptyFetch++
		}
	case upstream.FetchValid, upstream.FetchDuplicateOnly:
		stats.ConsecutiveEmptyFetch = 0
	case upstream.FetchError:
		if stats.FetchErrorCount < math.MaxInt64 {
			stats.FetchErrorCount++
		}
	default:
		return
	}
	recorder.byID[result.UpstreamID] = stats
}

func (recorder *StatsRecorder) ReadProviderStats(upstreamIDs []string) []Stats {
	result := make([]Stats, len(upstreamIDs))
	if recorder == nil {
		return result
	}
	recorder.mu.Lock()
	defer recorder.mu.Unlock()
	for index, upstreamID := range upstreamIDs {
		result[index] = recorder.byID[upstreamID]
		result[index].UpstreamID = upstreamID
	}
	return result
}

// RetainProviderStats removes observations for upstreams no longer present in
// the complete configuration. Disabled but configured upstreams must be kept.
func (recorder *StatsRecorder) RetainProviderStats(upstreamIDs []string) {
	if recorder == nil {
		return
	}
	retained := make(map[string]struct{}, len(upstreamIDs))
	for _, upstreamID := range upstreamIDs {
		if upstreamID != "" {
			retained[upstreamID] = struct{}{}
		}
	}
	recorder.mu.Lock()
	defer recorder.mu.Unlock()
	for upstreamID := range recorder.byID {
		if _, keep := retained[upstreamID]; !keep {
			delete(recorder.byID, upstreamID)
		}
	}
}