package service

import (
	"net"
	"sync"
	"time"
)

const defaultSourceIPRelease = 30 * time.Minute

type sourceIPEntry struct {
	connections int
	lastSeen    time.Time
}

type sourceIPBucket struct {
	mu      sync.Mutex
	limit   int
	release time.Duration
	entries map[string]*sourceIPEntry
}

type sourceIPLease struct {
	bucket *sourceIPBucket
	ip     string
	once   sync.Once
}

func (l *sourceIPLease) Release() {
	if l == nil || l.bucket == nil {
		return
	}
	l.once.Do(func() {
		l.bucket.mu.Lock()
		defer l.bucket.mu.Unlock()
		entry := l.bucket.entries[l.ip]
		if entry == nil {
			return
		}
		if entry.connections > 0 {
			entry.connections--
		}
		entry.lastSeen = time.Now()
		if entry.connections == 0 && l.bucket.release <= 0 {
			delete(l.bucket.entries, l.ip)
		}
	})
}

type sourceIPLimitManager struct {
	mu      sync.RWMutex
	buckets map[string]*sourceIPBucket
}

var globalSourceIPLimitManager = newSourceIPLimitManager()

func newSourceIPLimitManager() *sourceIPLimitManager {
	m := &sourceIPLimitManager{buckets: make(map[string]*sourceIPBucket)}
	go m.cleanupLoop()
	return m
}

func (m *sourceIPLimitManager) Configure(key string, limit int, release time.Duration) {
	if key == "" {
		return
	}
	if release < 0 {
		release = defaultSourceIPRelease
	}
	m.mu.Lock()
	bucket := m.buckets[key]
	if bucket == nil {
		bucket = &sourceIPBucket{entries: make(map[string]*sourceIPEntry)}
		m.buckets[key] = bucket
	}
	m.mu.Unlock()

	bucket.mu.Lock()
	bucket.limit = limit
	bucket.release = release
	bucket.mu.Unlock()
}

func (m *sourceIPLimitManager) Acquire(key, sourceIP string) (*sourceIPLease, bool) {
	if key == "" {
		return nil, true
	}
	m.mu.RLock()
	bucket := m.buckets[key]
	m.mu.RUnlock()
	if bucket == nil {
		return nil, true
	}

	ip := normalizeSourceIP(sourceIP)
	now := time.Now()
	bucket.mu.Lock()
	defer bucket.mu.Unlock()
	if bucket.limit <= 0 {
		return nil, true
	}
	bucket.cleanupLocked(now)
	if entry := bucket.entries[ip]; entry != nil {
		entry.connections++
		entry.lastSeen = now
		return &sourceIPLease{bucket: bucket, ip: ip}, true
	}
	if len(bucket.entries) >= bucket.limit {
		return nil, false
	}
	bucket.entries[ip] = &sourceIPEntry{connections: 1, lastSeen: now}
	return &sourceIPLease{bucket: bucket, ip: ip}, true
}

func (b *sourceIPBucket) cleanupLocked(now time.Time) {
	for ip, entry := range b.entries {
		if entry.connections == 0 && (b.release <= 0 || now.Sub(entry.lastSeen) >= b.release) {
			delete(b.entries, ip)
		}
	}
}

func (m *sourceIPLimitManager) cleanupLoop() {
	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()
	for now := range ticker.C {
		m.mu.RLock()
		buckets := make([]*sourceIPBucket, 0, len(m.buckets))
		for _, bucket := range m.buckets {
			buckets = append(buckets, bucket)
		}
		m.mu.RUnlock()
		for _, bucket := range buckets {
			bucket.mu.Lock()
			bucket.cleanupLocked(now)
			bucket.mu.Unlock()
		}
	}
}

func normalizeSourceIP(value string) string {
	if host, _, err := net.SplitHostPort(value); err == nil {
		value = host
	}
	if ip := net.ParseIP(value); ip != nil {
		return ip.String()
	}
	return value
}
