package lh

import (
	"bytes"
	"fmt"
	"sort"
	"sync"
	"time"

	"github.com/cespare/xxhash/v2"
	"github.com/olekukonko/ll/lx"
)

// Dedup is a log handler that suppresses duplicate entries within a TTL window.
// It wraps another handler and filters out repeated log entries that match
// within the deduplication period.
type Dedup struct {
	next           lx.Handler
	ttl            time.Duration
	cleanupEvery   time.Duration
	keyFn          lx.Deduper
	maxKeys        int
	excludedFields map[string]struct{} // Fields to exclude from default key hash
	shards         [32]dedupShard      // Shards reduce lock contention
	done           chan struct{}
	wg             sync.WaitGroup
	once           sync.Once
}

type dedupShard struct {
	mu   sync.Mutex
	seen map[uint64]int64 // key -> expiry timestamp
}

// DedupOpt configures a Dedup handler.
type DedupOpt func(*Dedup)

// WithDedupKeyFunc customizes how deduplication keys are generated.
func WithDedupKeyFunc(fn func(*lx.Entry) uint64) DedupOpt {
	return func(d *Dedup) {
		d.keyFn = dedupKeyFunc(fn)
	}
}

type dedupKeyFunc func(*lx.Entry) uint64

func (f dedupKeyFunc) Calculate(e *lx.Entry) uint64 {
	return f(e)
}

// WithDedupCleanupInterval sets how often expired deduplication keys are purged.
func WithDedupCleanupInterval(every time.Duration) DedupOpt {
	return func(d *Dedup) {
		if every > 0 {
			d.cleanupEvery = every
		}
	}
}

// WithDedupMaxKeys sets a soft limit on tracked deduplication keys.
func WithDedupMaxKeys(max int) DedupOpt {
	return func(d *Dedup) {
		if max > 0 {
			d.maxKeys = max
		}
	}
}

// WithDedupIgnore specifies fields to ignore in the default key function.
func WithDedupIgnore(fields ...string) DedupOpt {
	return func(d *Dedup) {
		if dd, ok := d.keyFn.(*defaultDedup); ok {
			if dd.ignoreFields == nil {
				dd.ignoreFields = make(map[string]struct{}, len(fields))
			}
			for _, f := range fields {
				dd.ignoreFields[f] = struct{}{}
			}
		}
	}
}

// NewDedup creates a deduplicating handler wrapper.
func NewDedup(next lx.Handler, ttl time.Duration, opts ...DedupOpt) *Dedup {
	if ttl <= 0 {
		ttl = 2 * time.Second
	}
	d := &Dedup{
		next:         next,
		ttl:          ttl,
		cleanupEvery: time.Minute,
		keyFn:        NewDefaultDedup(),
		done:         make(chan struct{}),
	}
	// Initialize shards with pre-allocated maps
	for i := 0; i < len(d.shards); i++ {
		d.shards[i].seen = make(map[uint64]int64, 64)
	}
	for _, opt := range opts {
		opt(d)
	}
	d.wg.Add(1)
	go d.cleanupLoop()
	return d
}

// Handle processes a log entry, suppressing duplicates within the TTL window.
func (d *Dedup) Handle(e *lx.Entry) error {
	now := time.Now().UnixNano()
	key := d.keyFn.Calculate(e)
	shardIdx := key % uint64(len(d.shards))
	shard := &d.shards[shardIdx]

	shard.mu.Lock()
	exp, ok := shard.seen[key]
	if ok && now < exp {
		shard.mu.Unlock()
		return nil
	}

	// Opportunistic cleanup if shard is getting full
	limitPerShard := d.maxKeys / len(d.shards)
	if d.maxKeys > 0 && len(shard.seen) >= limitPerShard {
		d.cleanupShardLocked(shard, now)
	}

	shard.seen[key] = now + d.ttl.Nanoseconds()
	shard.mu.Unlock()

	return d.next.Handle(e)
}

// Close stops the cleanup goroutine and closes the underlying handler.
func (d *Dedup) Close() error {
	var err error
	d.once.Do(func() {
		close(d.done)
		d.wg.Wait()
		if c, ok := d.next.(interface{ Close() error }); ok {
			err = c.Close()
		}
	})
	return err
}

// cleanupLoop runs periodically to purge expired deduplication keys.
func (d *Dedup) cleanupLoop() {
	defer d.wg.Done()
	ticker := time.NewTicker(d.cleanupEvery)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			now := time.Now().UnixNano()
			// Cleanup all shards sequentially to avoid CPU spike
			for i := 0; i < len(d.shards); i++ {
				shard := &d.shards[i]
				shard.mu.Lock()
				d.cleanupShardLocked(shard, now)
				shard.mu.Unlock()
			}
		case <-d.done:
			return
		}
	}
}

// cleanupShardLocked removes expired keys from a shard (caller must hold lock).
func (d *Dedup) cleanupShardLocked(shard *dedupShard, now int64) {
	for k, exp := range shard.seen {
		if now > exp {
			delete(shard.seen, k)
		}
	}
}

// defaultDedup implements the default deduplication key calculation.
type defaultDedup struct {
	ignoreFields map[string]struct{}
}

func NewDefaultDedup() lx.Deduper {
	return &defaultDedup{ignoreFields: make(map[string]struct{})}
}

// Calculate generates a deduplication key from level, message, namespace, and fields.
// Excludes any fields in ignoreFields. Uses xxhash for fast hashing.
func (d *defaultDedup) Calculate(e *lx.Entry) uint64 {
	h := xxhash.New()
	zero := []byte{0}

	// Core fields - use Write([]byte) for all data
	h.WriteString(e.Level.String())
	h.Write(zero)
	h.WriteString(e.Message)
	h.Write(zero)
	h.WriteString(e.Namespace)
	h.Write(zero)

	// Add fields if present
	if len(e.Fields) > 0 {
		m := e.Fields.Map()
		keys := make([]string, 0, len(m))
		for k := range m {
			if _, excluded := d.ignoreFields[k]; !excluded {
				keys = append(keys, k)
			}
		}
		sort.Strings(keys)

		// Use pooled buffer to avoid allocations
		buf := dedupBufPool.Get().(*bytes.Buffer)
		buf.Reset()
		defer dedupBufPool.Put(buf)

		for _, k := range keys {
			buf.WriteString(k)
			buf.WriteByte('=')
			fmt.Fprint(buf, m[k])
			buf.WriteByte(0)
		}
		h.Write(buf.Bytes())
	}

	return h.Sum64()
}

var dedupBufPool = sync.Pool{
	New: func() any { return new(bytes.Buffer) },
}
