package csmap

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

// CsMap 基于官方 sync.Map 重构，完全兼容原有接口
type CsMap[K comparable, V any] struct {
	inner sync.Map // 核心替换为官方 sync.Map
}

// 以下配置方法保留（兼容原有调用方式，但内部无实际作用）
type OptFunc[K comparable, V any] func(o *CsMap[K, V])

// New 创建基于 sync.Map 的并发安全 Map，兼容原有配置参数（参数无实际作用）
func New[K comparable, V any](options ...OptFunc[K, V]) *CsMap[K, V] {
	m := &CsMap[K, V]{}
	// 遍历配置项（兼容原有代码，无实际逻辑）
	for _, option := range options {
		option(m)
	}
	return m
}

// 保留原有配置方法（空实现，保证接口兼容）
func WithShardCount[K comparable, V any](count uint64) func(csMap *CsMap[K, V]) {
	return func(csMap *CsMap[K, V]) {}
}

func WithCustomHasher[K comparable, V any](h func(key K) uint64) func(csMap *CsMap[K, V]) {
	return func(csMap *CsMap[K, V]) {}
}

func WithSize[K comparable, V any](size uint64) func(csMap *CsMap[K, V]) {
	return func(csMap *CsMap[K, V]) {}
}

// -------------------------- 核心操作方法（基于 sync.Map 实现） --------------------------

// Store 存储键值对，兼容原有接口
func (m *CsMap[K, V]) Store(key K, value V) {
	m.inner.Store(key, value)
}

// Delete 删除指定键，返回是否删除成功
func (m *CsMap[K, V]) Delete(key K) bool {
	// sync.Map.Delete 无返回值，需先 Load 判断是否存在
	_, ok := m.inner.Load(key)
	if ok {
		m.inner.Delete(key)
	}
	return ok
}

// DeleteIf 满足条件时删除
func (m *CsMap[K, V]) DeleteIf(key K, condition func(value V) bool) bool {
	// 先 Load 获取值，再判断条件
	val, ok := m.inner.Load(key)
	if !ok {
		return false
	}

	v, okCast := val.(V)
	if !okCast {
		return false
	}

	if condition(v) {
		m.inner.Delete(key)
		return true
	}
	return false
}

// Load 获取指定键的值
func (m *CsMap[K, V]) Load(key K) (V, bool) {
	var zero V
	val, ok := m.inner.Load(key)
	if !ok {
		return zero, false
	}

	// 类型断言（保证类型安全）
	v, okCast := val.(V)
	if !okCast {
		return zero, false
	}
	return v, true
}

// Has 判断键是否存在
func (m *CsMap[K, V]) Has(key K) bool {
	_, ok := m.inner.Load(key)
	return ok
}

// Clear 清空所有数据
func (m *CsMap[K, V]) Clear() {
	// sync.Map 无直接 Clear 方法，通过 Range 遍历删除
	m.inner.Range(func(key, value any) bool {
		m.inner.Delete(key)
		return true
	})
}

// Count 统计元素数量
func (m *CsMap[K, V]) Count() int {
	count := 0
	m.inner.Range(func(key, value any) bool {
		count++
		return true
	})
	return count
}

// SetIfAbsent 仅当键不存在时设置值
func (m *CsMap[K, V]) SetIfAbsent(key K, value V) {
	m.inner.LoadOrStore(key, value)
}

// SetIf 根据条件设置值
func (m *CsMap[K, V]) SetIf(key K, conditionFn func(previousValue V, previousFound bool) (value V, set bool)) {
	prevVal, found := m.inner.Load(key)
	var prevV V
	if found {
		prevV, _ = prevVal.(V)
	}

	// 执行条件函数
	newVal, set := conditionFn(prevV, found)
	if set {
		m.inner.Store(key, newVal)
	}
}

// SetIfPresent 仅当键存在时设置值
func (m *CsMap[K, V]) SetIfPresent(key K, value V) {
	// 先判断是否存在，再设置
	if _, ok := m.inner.Load(key); ok {
		m.inner.Store(key, value)
	}
}

// IsEmpty 判断是否为空
func (m *CsMap[K, V]) IsEmpty() bool {
	return m.Count() == 0
}

// Tuple 保留原有结构体（兼容序列化逻辑）
type Tuple[K comparable, V any] struct {
	Key K
	Val V
}

// -------------------------- 关键修复：Range 方法（无锁阻塞风险） --------------------------
func (m *CsMap[K, V]) Range(f func(key K, value V) (stop bool)) {
	if f == nil {
		return
	}

	var stopFlag atomic.Bool

	// 基于 sync.Map 的 Range 实现，无额外 goroutine/channel
	m.inner.Range(func(key, value any) bool {
		// 检测终止标志
		if stopFlag.Load() {
			return false
		}

		// 类型断言
		k, okK := key.(K)
		v, okV := value.(V)
		if !okK || !okV {
			return true // 类型不匹配时跳过，继续遍历
		}

		// 执行用户回调
		if f(k, v) {
			stopFlag.Store(true)
			return false // 终止遍历
		}
		return true
	})
}

// -------------------------- 序列化方法（兼容原有逻辑） --------------------------
func (m *CsMap[K, V]) MarshalJSON() ([]byte, error) {
	tmp := make(map[K]V, m.Count())
	m.Range(func(key K, value V) (stop bool) {
		tmp[key] = value
		return false
	})
	return json.Marshal(tmp)
}

func (m *CsMap[K, V]) UnmarshalJSON(b []byte) error {
	tmp := make(map[K]V)
	if err := json.Unmarshal(b, &tmp); err != nil {
		return err
	}

	// 清空原有数据
	m.Clear()
	// 批量存储
	for key, val := range tmp {
		m.Store(key, val)
	}
	return nil
}

// -------------------------- 移除所有无用的旧方法（produce/listen 等） --------------------------