package utils

import (
	"encoding/json"
	"sort"
	"sync"
)

// OrderMap 是一个支持排序、有序遍历、并发安全的泛型 map。
// 遍历（Range / Iter / IterBuffered）为非阻塞：进入遍历时创建快照（短时间持锁），
// 然后在快照上遍历（不持锁），因此回调可以做耗时操作而不阻塞写操作。
type OrderMap[K comparable, V any] struct {
	mu   sync.RWMutex
	data map[K]V
	keys []K
	less func(a, b K) bool // key 排序规则；为 nil 时按插入顺序
}

// NewOrderedMap 创建一个 OrderMap；less 为 nil 时按插入顺序，否则会在写入后按 less 排序 keys。
func NewOrderedMap[K comparable, V any](less func(a, b K) bool) *OrderMap[K, V] {
	return &OrderMap[K, V]{
		data: make(map[K]V),
		keys: make([]K, 0),
		less: less,
	}
}

// Set 设置键值（存在则覆盖，不存在则添加）。若提供了 less，则会在插入后排序 keys。
func (om *OrderMap[K, V]) Set(k K, v V) {
	om.mu.Lock()
	if _, exist := om.data[k]; !exist {
		om.keys = append(om.keys, k)
	}
	om.data[k] = v
	// 若定义了 less，则保持 keys 有序
	if om.less != nil {
		sort.Slice(om.keys, func(i, j int) bool { return om.less(om.keys[i], om.keys[j]) })
	}
	om.mu.Unlock()
}

// Get 获取一个键的值，返回值和是否存在（并发安全）。
func (om *OrderMap[K, V]) Get(k K) (V, bool) {
	om.mu.RLock()
	v, ok := om.data[k]
	om.mu.RUnlock()
	return v, ok
}

// Delete 删除一个键（并发安全）。
func (om *OrderMap[K, V]) Delete(k K) {
	om.mu.Lock()
	if _, ok := om.data[k]; ok {
		delete(om.data, k)
		// 从 keys 切片中移除第一次出现的 k
		for i, key := range om.keys {
			if key == k {
				om.keys = append(om.keys[:i], om.keys[i+1:]...)
				break
			}
		}
	}
	om.mu.Unlock()
}

// Keys 返回当前 keys 的有序副本（并发安全）。
func (om *OrderMap[K, V]) Keys() []K {
	om.mu.RLock()
	keys := make([]K, len(om.keys))
	copy(keys, om.keys)
	om.mu.RUnlock()
	return keys
}

// Values 返回当前 values 的有序副本（并发安全）。
func (om *OrderMap[K, V]) Values() []V {
	om.mu.RLock()
	values := make([]V, 0, len(om.keys))
	for _, k := range om.keys {
		values = append(values, om.data[k])
	}
	om.mu.RUnlock()
	return values
}

// Len 返回 map 长度（并发安全）。
func (om *OrderMap[K, V]) Len() int {
	om.mu.RLock()
	n := len(om.data)
	om.mu.RUnlock()
	return n
}

// snapshot 复制一份 keys 列表和对应的数据（短时间持锁），返回 keys 及 map 的副本。
// 注意：这是浅拷贝（V 为引用类型则仍为相同引用）。
func (om *OrderMap[K, V]) snapshot() ([]K, map[K]V) {
	om.mu.RLock()
	keys := make([]K, len(om.keys))
	copy(keys, om.keys)
	dataCopy := make(map[K]V, len(keys))
	for _, k := range keys {
		dataCopy[k] = om.data[k]
	}
	om.mu.RUnlock()
	return keys, dataCopy
}

// Iter 返回一个无缓冲的 channel，用于在 for range 中遍历所有当前项。
// Deprecated: 使用 IterBuffered() 性能更好。
func (om *OrderMap[K, V]) Iter() <-chan Tuple[K, V] {
	keys, dataCopy := om.snapshot()
	ch := make(chan Tuple[K, V])
	go func() {
		for _, k := range keys {
			// 可能在原 map 中此 key 已被删除，但 snapshot 中仍存在
			if v, ok := dataCopy[k]; ok {
				ch <- Tuple[K, V]{Key: k, Val: v}
			}
		}
		close(ch)
	}()
	return ch
}

// IterBuffered 返回一个带缓冲的 channel（缓冲大小为当前元素数），用于高效遍历。
func (om *OrderMap[K, V]) IterBuffered() <-chan Tuple[K, V] {
	keys, dataCopy := om.snapshot()
	ch := make(chan Tuple[K, V], len(keys))
	go func() {
		for _, k := range keys {
			if v, ok := dataCopy[k]; ok {
				ch <- Tuple[K, V]{Key: k, Val: v}
			}
		}
		close(ch)
	}()
	return ch
}

// Range 在非阻塞快照上按有序 keys 进行遍历。f 返回 false 可提前停止遍历。
// Range 本身不会在回调期间持锁，因此回调可以做耗时操作。
func (om *OrderMap[K, V]) Range(f func(K, V) bool) {
	keys, dataCopy := om.snapshot()
	for _, k := range keys {
		v, ok := dataCopy[k]
		if !ok {
			continue
		}
		if !f(k, v) {
			break
		}
	}
}

// Items 返回一份 map 的浅拷贝（并发安全）。注意：返回 map 的 key 顺序不可保证（map 本身无序）
// 若你需要有序的键值序列，使用 IterBuffered/Range/Keys。
func (om *OrderMap[K, V]) Items() map[K]V {
	_, dataCopy := om.snapshot()
	// snapshot 已是 map 的拷贝，直接返回
	return dataCopy
}

// Clear 清空所有数据（并发安全）。
func (om *OrderMap[K, V]) Clear() {
	om.mu.Lock()
	om.data = make(map[K]V)
	om.keys = make([]K, 0)
	om.mu.Unlock()
}

// ----------------------------------------------------------------------------
// JSON 支持（保持有序）
// ----------------------------------------------------------------------------

// jsonPair 用于 Marshal/Unmarshal 时保留顺序的中间结构：数组形式 [{key:..., value:...}, ...]
type jsonPair[K comparable, V any] struct {
	Key   K `json:"key"`
	Value V `json:"value"`
}

// MarshalJSON 将 OrderMap 以数组形式序列化，保持 keys 的顺序：
// [
//
//	{"key": <k1>, "value": <v1>},
//	{"key": <k2>, "value": <v2>},
//	...
//
// ]
func (om *OrderMap[K, V]) MarshalJSON() ([]byte, error) {
	keys, dataCopy := om.snapshot()
	pairs := make([]jsonPair[K, V], 0, len(keys))
	for _, k := range keys {
		if v, ok := dataCopy[k]; ok {
			pairs = append(pairs, jsonPair[K, V]{Key: k, Value: v})
		}
	}
	return json.Marshal(pairs)
}

// UnmarshalJSON 从上述数组格式反序列化，按数组顺序恢复 keys 与数据。
// 注意：若 Key 类型不是 JSON 可直接解码的类型，反序列化会失败。
func (om *OrderMap[K, V]) UnmarshalJSON(b []byte) error {
	var pairs []jsonPair[K, V]
	if err := json.Unmarshal(b, &pairs); err != nil {
		return err
	}
	om.mu.Lock()
	defer om.mu.Unlock()
	om.data = make(map[K]V, len(pairs))
	om.keys = make([]K, 0, len(pairs))
	for _, p := range pairs {
		om.data[p.Key] = p.Value
		om.keys = append(om.keys, p.Key)
	}
	// 若定义了 less，则保持 keys 排序
	if om.less != nil {
		sort.Slice(om.keys, func(i, j int) bool { return om.less(om.keys[i], om.keys[j]) })
	}
	return nil
}