package activitypool

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"sort"
	"sync"
	"time"

	extractionDomain "proxy-pool/internal/domain/extraction"
	ownershipDomain "proxy-pool/internal/domain/ownership"
	proxyDomain "proxy-pool/internal/domain/proxy"
)

const defaultIdempotencyTTL = 5 * time.Minute

var ErrInvalidBatch = errors.New("invalid activity pool batch")

// FetchedBatch describes one ephemeral provider response. Proxies without a
// usable expiry are dropped because this pool is intentionally rebuildable.
type FetchedBatch struct {
	ObservedAt             time.Time
	ConfiguredTTL          time.Duration
	AllocationSafetyMargin time.Duration
	Proxies                []proxyDomain.Proxy
}

type UpsertResult struct {
	Accepted  int
	Inserted  int
	Refreshed int
	Dropped   int
}

type Upserter interface {
	// A proxy unique key has one incumbent upstream for its current lifetime.
	// Duplicates from other upstreams must not replace that source lifecycle.
	UpsertFetched(context.Context, string, FetchedBatch) (UpsertResult, error)
}

type Entry struct {
	Proxy         proxyDomain.Proxy
	UsableUntil   time.Time
	OwnerWorkerID string
	State         proxyDomain.State
}

type MemoryPool struct {
	mu sync.Mutex

	entries    map[string]Entry
	keyByID    map[string]string
	idempotent map[string]idempotencyEntry
	ownership  map[string]ownershipDomain.Assignment
	nextEpoch  uint64
}

type idempotencyEntry struct {
	command   extractionDomain.Command
	result    extractionDomain.Result
	expiresAt time.Time
}

var (
	_ Upserter                   = (*MemoryPool)(nil)
	_ extractionDomain.Store     = (*MemoryPool)(nil)
	_ ownershipDomain.Repository = (*MemoryPool)(nil)
)

func NewMemoryPool() *MemoryPool {
	return &MemoryPool{
		entries:    make(map[string]Entry),
		keyByID:    make(map[string]string),
		idempotent: make(map[string]idempotencyEntry),
		ownership:  make(map[string]ownershipDomain.Assignment),
	}
}

func (p *MemoryPool) UpsertFetched(ctx context.Context, upstreamID string, batch FetchedBatch) (UpsertResult, error) {
	var result UpsertResult
	if err := ctx.Err(); err != nil {
		return result, err
	}
	if p == nil || upstreamID == "" || batch.ObservedAt.IsZero() || batch.ConfiguredTTL < 0 ||
		batch.AllocationSafetyMargin < 0 ||
		(batch.ConfiguredTTL > 0 && batch.AllocationSafetyMargin >= batch.ConfiguredTTL) {
		return result, ErrInvalidBatch
	}
	for _, candidate := range batch.Proxies {
		if candidate.SourceUpstream != "" && candidate.SourceUpstream != upstreamID {
			return result, ErrInvalidBatch
		}
	}

	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return result, err
	}
	p.purgeExpiredLocked(batch.ObservedAt)
	seenIDs := make(map[string]string, len(batch.Proxies))
	for _, candidate := range batch.Proxies {
		if !validProxyIdentity(candidate) {
			continue
		}
		key := candidate.UniqueKey()
		proxyID := candidate.ID
		if proxyID == "" {
			proxyID = stableProxyID(key)
		}
		if existingKey, exists := p.keyByID[proxyID]; exists && existingKey != key {
			return UpsertResult{}, ErrInvalidBatch
		}
		if existingKey, exists := seenIDs[proxyID]; exists && existingKey != key {
			return UpsertResult{}, ErrInvalidBatch
		}
		seenIDs[proxyID] = key
	}

	for _, candidate := range batch.Proxies {
		if !validProxyIdentity(candidate) {
			result.Dropped++
			continue
		}
		candidate.SourceUpstream = upstreamID
		expiresAt := proxyDomain.EffectiveExpiry(batch.ObservedAt, candidate.ExpiresAt, 0, batch.ConfiguredTTL)
		if expiresAt == nil {
			result.Dropped++
			continue
		}
		usableUntil := expiresAt.Add(-batch.AllocationSafetyMargin)
		if !usableUntil.After(batch.ObservedAt) {
			result.Dropped++
			continue
		}
		candidate.ExpiresAt = expiresAt
		candidate.UsableUntil = &usableUntil
		if candidate.CreatedAt.IsZero() {
			candidate.CreatedAt = batch.ObservedAt.UTC()
		}
		if candidate.State == "" {
			candidate.State = proxyDomain.StateFetched
		}
		key := candidate.UniqueKey()
		result.Accepted++

		if current, exists := p.entries[key]; exists {
			result.Refreshed++
			if current.State == proxyDomain.StateExtracted || current.Proxy.SourceUpstream != upstreamID {
				continue
			}
			candidate.ID = current.Proxy.ID
			candidate.CreatedAt = current.Proxy.CreatedAt
			candidate.State = current.State
			candidate.LastCheckedAt = current.Proxy.LastCheckedAt
			candidate.LastSuccessAt = current.Proxy.LastSuccessAt
			candidate.Latency = current.Proxy.Latency
			p.entries[key] = Entry{
				Proxy:         cloneProxy(candidate),
				UsableUntil:   usableUntil,
				OwnerWorkerID: current.OwnerWorkerID,
				State:         current.State,
			}
			continue
		}

		if candidate.ID == "" {
			candidate.ID = stableProxyID(key)
		}
		p.entries[key] = Entry{
			Proxy:       cloneProxy(candidate),
			UsableUntil: usableUntil,
			State:       candidate.State,
		}
		p.keyByID[candidate.ID] = key
		result.Inserted++
	}
	return result, nil
}

func (p *MemoryPool) Snapshot(now time.Time) []Entry {
	if p == nil {
		return nil
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	p.purgeExpiredLocked(now)
	entries := make([]Entry, 0, len(p.entries))
	for _, entry := range p.entries {
		entry.Proxy = cloneProxy(entry.Proxy)
		entries = append(entries, entry)
	}
	sort.Slice(entries, func(i, j int) bool { return entries[i].Proxy.ID < entries[j].Proxy.ID })
	return entries
}

func (p *MemoryPool) PurgeExpired(now time.Time) int {
	if p == nil {
		return 0
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.purgeExpiredLocked(now)
}

func (p *MemoryPool) Extract(ctx context.Context, command extractionDomain.Command) (extractionDomain.Result, error) {
	result := extractionDomain.Result{Requested: command.Requested}
	if err := ctx.Err(); err != nil {
		return result, err
	}
	if p == nil || command.Now.IsZero() || command.Requested < 0 || command.ReserveForGateway < 0 ||
		command.MinRemainingTTL < 0 || command.MaxHealthCheckAge < 0 || command.IdempotencyTTL < 0 ||
		(command.IdempotencyKey != "" && command.ClientID == "") ||
		(command.Fulfillment != extractionDomain.Partial && command.Fulfillment != extractionDomain.AllOrNothing) {
		return result, extractionDomain.ErrInvalidCommand
	}

	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return result, err
	}
	p.purgeExpiredLocked(command.Now)

	idempotencyKey := command.ClientID + "\x00" + command.IdempotencyKey
	if command.IdempotencyKey != "" {
		if committed, ok := p.idempotent[idempotencyKey]; ok {
			if !sameIdempotentRequest(committed.command, command) {
				return result, extractionDomain.ErrIdempotencyConflict
			}
			return cloneResult(committed.result), nil
		}
	}
	if command.Requested == 0 {
		return result, nil
	}

	eligible := make([]string, 0, len(p.entries))
	for key, entry := range p.entries {
		if eligibleForExtraction(entry, command) {
			eligible = append(eligible, key)
		}
	}
	sort.Slice(eligible, func(i, j int) bool {
		return p.entries[eligible[i]].UsableUntil.After(p.entries[eligible[j]].UsableUntil)
	})
	available := len(eligible) - command.ReserveForGateway
	if available < 0 {
		available = 0
	}
	if command.Fulfillment == extractionDomain.AllOrNothing && available < command.Requested {
		return result, extractionDomain.ErrInsufficientProxies
	}
	count := command.Requested
	if count > available {
		count = available
	}
	for _, key := range eligible[:count] {
		entry := p.entries[key]
		entry.State = proxyDomain.StateExtracted
		entry.Proxy.State = proxyDomain.StateExtracted
		p.entries[key] = entry
		result.Items = append(result.Items, extractionCandidate(entry))
	}
	result.Returned = len(result.Items)
	if result.Returned > 0 {
		result.ExtractedAt = command.Now.UTC()
	}
	if command.IdempotencyKey != "" {
		expiresAt := command.Now.Add(idempotencyTTL(command.IdempotencyTTL))
		for _, item := range result.Items {
			if !item.ExpiresAt.IsZero() && item.ExpiresAt.Before(expiresAt) {
				expiresAt = item.ExpiresAt
			}
		}
		if expiresAt.After(command.Now) {
			p.idempotent[idempotencyKey] = idempotencyEntry{
				command: cloneCommand(command), result: cloneResult(result), expiresAt: expiresAt,
			}
		}
	}
	return result, nil
}

func (p *MemoryPool) Assign(ctx context.Context, now time.Time, proxyID, workerID string, ttl time.Duration) (ownershipDomain.Assignment, error) {
	if err := ownershipContextError(ctx); err != nil {
		return ownershipDomain.Assignment{}, err
	}
	if p == nil || now.IsZero() || proxyID == "" || workerID == "" || ttl <= 0 {
		return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return ownershipDomain.Assignment{}, err
	}
	p.purgeExpiredLocked(now)
	if current, exists := p.ownership[proxyID]; exists {
		if current.ExpiresAt.After(now) {
			return ownershipDomain.Assignment{}, ownershipDomain.ErrAlreadyOwned
		}
		if entry, ok := p.entryByIDLocked(proxyID); ok && entry.OwnerWorkerID == current.WorkerID {
			entry.OwnerWorkerID = ""
			p.setEntryByIDLocked(proxyID, entry)
		}
		delete(p.ownership, proxyID)
	}
	entry, ok := p.entryByIDLocked(proxyID)
	if !ok || entry.State != proxyDomain.StateAvailable || entry.OwnerWorkerID != "" || !entry.UsableUntil.After(now) {
		return ownershipDomain.Assignment{}, ownershipDomain.ErrOwnershipUnavailable
	}
	p.nextEpoch++
	expiresAt := minTime(now.UTC().Add(ttl), entry.UsableUntil)
	assignment := ownershipDomain.Assignment{
		ProxyID: proxyID, WorkerID: workerID, Epoch: p.nextEpoch, Version: 1, ExpiresAt: expiresAt,
	}
	entry.OwnerWorkerID = workerID
	p.setEntryByIDLocked(proxyID, entry)
	p.ownership[proxyID] = assignment
	return assignment, nil
}

func (p *MemoryPool) Renew(ctx context.Context, now time.Time, proxyID, workerID string, epoch uint64, ttl time.Duration) (ownershipDomain.Assignment, error) {
	if err := ownershipContextError(ctx); err != nil {
		return ownershipDomain.Assignment{}, err
	}
	if p == nil || now.IsZero() || ttl <= 0 {
		return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return ownershipDomain.Assignment{}, err
	}
	p.purgeExpiredLocked(now)
	assignment, ok := p.ownership[proxyID]
	if !ok || assignment.WorkerID != workerID || assignment.Epoch != epoch || !assignment.ExpiresAt.After(now) {
		return ownershipDomain.Assignment{}, ownershipDomain.ErrStaleAssignment
	}
	entry, ok := p.entryByIDLocked(proxyID)
	if !ok {
		return ownershipDomain.Assignment{}, ownershipDomain.ErrStaleAssignment
	}
	assignment.ExpiresAt = minTime(now.UTC().Add(ttl), entry.UsableUntil)
	assignment.Version++
	p.ownership[proxyID] = assignment
	return assignment, nil
}

func (p *MemoryPool) BeginDrain(ctx context.Context, proxyID, workerID string, epoch uint64) (ownershipDomain.Assignment, error) {
	if err := ownershipContextError(ctx); err != nil {
		return ownershipDomain.Assignment{}, err
	}
	if p == nil {
		return ownershipDomain.Assignment{}, ownershipDomain.ErrStaleAssignment
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return ownershipDomain.Assignment{}, err
	}
	assignment, ok := p.ownership[proxyID]
	if !ok || assignment.WorkerID != workerID || assignment.Epoch != epoch {
		return ownershipDomain.Assignment{}, ownershipDomain.ErrStaleAssignment
	}
	if !assignment.Draining {
		assignment.Draining = true
		assignment.Version++
		p.ownership[proxyID] = assignment
	}
	return assignment, nil
}

func (p *MemoryPool) AcknowledgeDrain(ctx context.Context, proxyID, workerID string, epoch uint64, active, reserved int64) error {
	if err := ownershipContextError(ctx); err != nil {
		return err
	}
	if p == nil || active < 0 || reserved < 0 {
		return ownershipDomain.ErrInvalidOwnership
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return err
	}
	assignment, ok := p.ownership[proxyID]
	if !ok || assignment.WorkerID != workerID || assignment.Epoch != epoch {
		return ownershipDomain.ErrStaleAssignment
	}
	if !assignment.Draining {
		return ownershipDomain.ErrNotDraining
	}
	if active > 0 || reserved > 0 {
		return ownershipDomain.ErrDrainNotReady
	}
	if entry, exists := p.entryByIDLocked(proxyID); exists && entry.OwnerWorkerID == workerID {
		entry.OwnerWorkerID = ""
		p.setEntryByIDLocked(proxyID, entry)
	}
	delete(p.ownership, proxyID)
	return nil
}

func (p *MemoryPool) Get(ctx context.Context, proxyID string) (ownershipDomain.Assignment, bool, error) {
	if err := ownershipContextError(ctx); err != nil {
		return ownershipDomain.Assignment{}, false, err
	}
	if p == nil {
		return ownershipDomain.Assignment{}, false, nil
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return ownershipDomain.Assignment{}, false, err
	}
	assignment, ok := p.ownership[proxyID]
	return assignment, ok, nil
}

func (p *MemoryPool) Expire(ctx context.Context, now time.Time, limit int) ([]ownershipDomain.Assignment, error) {
	if err := ownershipContextError(ctx); err != nil {
		return nil, err
	}
	if limit <= 0 {
		return nil, ownershipDomain.ErrInvalidOwnership
	}
	if p == nil {
		return nil, nil
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	eligibleIDs := make([]string, 0)
	for proxyID, assignment := range p.ownership {
		_, exists := p.entryByIDLocked(proxyID)
		if exists && assignment.ExpiresAt.After(now) {
			continue
		}
		eligibleIDs = append(eligibleIDs, proxyID)
	}
	sort.Strings(eligibleIDs)
	if len(eligibleIDs) > limit {
		eligibleIDs = eligibleIDs[:limit]
	}
	expired := make([]ownershipDomain.Assignment, 0, len(eligibleIDs))
	for _, proxyID := range eligibleIDs {
		assignment := p.ownership[proxyID]
		entry, exists := p.entryByIDLocked(proxyID)
		if exists && entry.OwnerWorkerID == assignment.WorkerID {
			entry.OwnerWorkerID = ""
			p.setEntryByIDLocked(proxyID, entry)
		}
		expired = append(expired, assignment)
		delete(p.ownership, proxyID)
	}
	return expired, nil
}

func ownershipContextError(ctx context.Context) error {
	if ctx == nil {
		return ownershipDomain.ErrInvalidOwnership
	}
	return ctx.Err()
}

func (p *MemoryPool) purgeExpiredLocked(now time.Time) int {
	removed := 0
	for key, entry := range p.entries {
		if entry.Proxy.ExpiresAt == nil || entry.Proxy.ExpiresAt.After(now) {
			continue
		}
		delete(p.ownership, entry.Proxy.ID)
		delete(p.keyByID, entry.Proxy.ID)
		delete(p.entries, key)
		removed++
	}
	for key, entry := range p.idempotent {
		if !entry.expiresAt.After(now) {
			delete(p.idempotent, key)
		}
	}
	return removed
}

func (p *MemoryPool) entryByIDLocked(proxyID string) (Entry, bool) {
	key, ok := p.keyByID[proxyID]
	if !ok {
		return Entry{}, false
	}
	entry, ok := p.entries[key]
	return entry, ok
}

func (p *MemoryPool) setEntryByIDLocked(proxyID string, entry Entry) {
	if key, ok := p.keyByID[proxyID]; ok {
		p.entries[key] = entry
	}
}

func eligibleForExtraction(entry Entry, command extractionDomain.Command) bool {
	if entry.State != proxyDomain.StateAvailable || entry.OwnerWorkerID != "" || !entry.UsableUntil.After(command.Now) {
		return false
	}
	if entry.Proxy.ExpiresAt != nil && entry.Proxy.ExpiresAt.Sub(command.Now) < command.MinRemainingTTL {
		return false
	}
	if command.MaxHealthCheckAge > 0 {
		if entry.Proxy.LastCheckedAt == nil || command.Now.Sub(*entry.Proxy.LastCheckedAt) > command.MaxHealthCheckAge {
			return false
		}
	}
	return matches(command.Protocols, string(entry.Proxy.Scheme)) &&
		matches(command.Regions, entry.Proxy.Tags["region"]) &&
		matches(command.Carriers, entry.Proxy.Tags["carrier"]) &&
		matches(command.Upstreams, entry.Proxy.SourceUpstream)
}

func extractionCandidate(entry Entry) extractionDomain.Candidate {
	expiresAt := time.Time{}
	if entry.Proxy.ExpiresAt != nil {
		expiresAt = *entry.Proxy.ExpiresAt
	}
	checkedAt := time.Time{}
	if entry.Proxy.LastCheckedAt != nil {
		checkedAt = *entry.Proxy.LastCheckedAt
	}
	return extractionDomain.Candidate{
		ID: entry.Proxy.ID, Protocol: string(entry.Proxy.Scheme), Host: entry.Proxy.Host,
		Port: entry.Proxy.Port, Username: entry.Proxy.Username, Region: entry.Proxy.Tags["region"],
		Carrier: entry.Proxy.Tags["carrier"], Upstream: entry.Proxy.SourceUpstream,
		OwnerWorkerID: entry.OwnerWorkerID, State: extractionDomain.Available,
		ExpiresAt: expiresAt, LastCheckedAt: checkedAt,
	}
}

func stableProxyID(key string) string {
	digest := sha256.Sum256([]byte(key))
	return "px_" + hex.EncodeToString(digest[:12])
}

func validProxyIdentity(candidate proxyDomain.Proxy) bool {
	if candidate.Host == "" || candidate.Port == 0 {
		return false
	}
	switch candidate.Scheme {
	case proxyDomain.SchemeHTTP, proxyDomain.SchemeHTTPS, proxyDomain.SchemeSOCKS5:
		return true
	default:
		return false
	}
}

func cloneProxy(candidate proxyDomain.Proxy) proxyDomain.Proxy {
	if candidate.ExpiresAt != nil {
		value := *candidate.ExpiresAt
		candidate.ExpiresAt = &value
	}
	if candidate.UsableUntil != nil {
		value := *candidate.UsableUntil
		candidate.UsableUntil = &value
	}
	if candidate.LastCheckedAt != nil {
		value := *candidate.LastCheckedAt
		candidate.LastCheckedAt = &value
	}
	if candidate.LastSuccessAt != nil {
		value := *candidate.LastSuccessAt
		candidate.LastSuccessAt = &value
	}
	if candidate.Tags != nil {
		tags := make(map[string]string, len(candidate.Tags))
		for key, value := range candidate.Tags {
			tags[key] = value
		}
		candidate.Tags = tags
	}
	return candidate
}

func cloneResult(result extractionDomain.Result) extractionDomain.Result {
	result.Items = append([]extractionDomain.Candidate(nil), result.Items...)
	return result
}

func cloneCommand(command extractionDomain.Command) extractionDomain.Command {
	command.Protocols = append([]string(nil), command.Protocols...)
	command.Regions = append([]string(nil), command.Regions...)
	command.Carriers = append([]string(nil), command.Carriers...)
	command.Upstreams = append([]string(nil), command.Upstreams...)
	return command
}

func sameIdempotentRequest(left, right extractionDomain.Command) bool {
	return left.Requested == right.Requested && left.Fulfillment == right.Fulfillment &&
		equalSet(left.Protocols, right.Protocols) && equalSet(left.Regions, right.Regions) &&
		equalSet(left.Carriers, right.Carriers) && equalSet(left.Upstreams, right.Upstreams)
}

func equalSet(left, right []string) bool {
	if len(left) != len(right) {
		return false
	}
	counts := make(map[string]int, len(left))
	for _, value := range left {
		counts[value]++
	}
	for _, value := range right {
		counts[value]--
		if counts[value] < 0 {
			return false
		}
	}
	return true
}

func matches(allowed []string, value string) bool {
	if len(allowed) == 0 {
		return true
	}
	for _, candidate := range allowed {
		if candidate == value {
			return true
		}
	}
	return false
}

func idempotencyTTL(configured time.Duration) time.Duration {
	if configured > 0 {
		return configured
	}
	return defaultIdempotencyTTL
}

func minTime(left, right time.Time) time.Time {
	if left.Before(right) {
		return left
	}
	return right
}