package kvstore

import (
    "context"
    "encoding/binary"
    "github.com/go-redis/redis/v8"
    "sync"
    "time"
)

// KVEntry 表示单个键值对缓存条目
type KVEntry struct {
    LayerID   int     // 模型层编号
    HeadID    int     // 注意力头编号
    Position  int     // 在序列中的位置
    KeyData   []float16 // 键向量
    ValueData []float16 // 值向量
    Score     float32 // 重要性分数，用于淘汰策略
    Timestamp int64   // 创建时间戳
}

// DistributedKVCache 分布式KV缓存管理器
type DistributedKVCache struct {
    redisClients []*redis.Client // Redis集群连接池
    localCache   *sync.Map       // 本地热缓存
    config       CacheConfig
}

type CacheConfig struct {
    RedisAddrs     []string // Redis地址列表
    LocalCacheSize int      // 本地缓存大小（条目数）
    Compression    bool     // 是否启用量化压缩
    QuantBits      int      // 量化位数，如8或4
}

// NewDistributedKVCache 创建分布式缓存实例
func NewDistributedKVCache(config CacheConfig) *DistributedKVCache {
    clients := make([]*redis.Client, len(config.RedisAddrs))
    for i, addr := range config.RedisAddrs {
        clients[i] = redis.NewClient(&redis.Options{
            Addr: addr,
            // 连接池配置优化长连接
            PoolSize:     100,
            MinIdleConns: 20,
        })
    }
    
    return &DistributedKVCache{
        redisClients: clients,
        localCache:   &sync.Map{},
        config:       config,
    }
}

// StoreKV 存储KV缓存到分布式系统
// 使用一致性哈希选择存储节点
func (c *DistributedKVCache) StoreKV(ctx context.Context, entry *KVEntry) error {
    // 1. 对KV数据进行量化压缩（如果启用）
    compressedKey, compressedValue := entry.KeyData, entry.ValueData
    if c.config.Compression {
        compressedKey = quantize(entry.KeyData, c.config.QuantBits)
        compressedValue = quantize(entry.ValueData, c.config.QuantBits)
    }
    
    // 2. 生成唯一键
    cacheKey := generateCacheKey(entry.LayerID, entry.HeadID, entry.Position)
    
    // 3. 序列化数据
    data, err := serializeEntry(entry, compressedKey, compressedValue)
    if err != nil {
        return err
    }
    
    // 4. 一致性哈希选择Redis节点
    nodeIndex := hash(cacheKey) % len(c.redisClients)
    
    // 5. 异步写入Redis，设置过期时间防止无限增长
    pipe := c.redisClients[nodeIndex].Pipeline()
    pipe.Set(ctx, cacheKey, data, 30*time.Minute)
    
    // 6. 同时更新本地热缓存
    c.localCache.Store(cacheKey, entry)
    
    _, err = pipe.Exec(ctx)
    return err
}

// BatchLoadKV 批量加载KV缓存，优化长序列访问
func (c *DistributedKVCache) BatchLoadKV(ctx context.Context, 
    layerID int, headID int, startPos, endPos int) ([]*KVEntry, error) {
    
    // 构建批量键
    keys := make([]string, 0, endPos-startPos+1)
    for pos := startPos; pos <= endPos; pos++ {
        keys = append(keys, generateCacheKey(layerID, headID, pos))
    }
    
    // 按Redis节点分组，减少网络往返
    nodeKeys := make(map[int][]string)
    for _, key := range keys {
        nodeIndex := hash(key) % len(c.redisClients)
        nodeKeys[nodeIndex] = append(nodeKeys[nodeIndex], key)
    }
    
    // 并发从各节点加载数据
    results := make([]*KVEntry, 0, len(keys))
    var mu sync.Mutex
    var wg sync.WaitGroup
    
    for nodeIdx, nodeKeys := range nodeKeys {
        wg.Add(1)
        go func(idx int, ks []string) {
            defer wg.Done()
            
            // 先查本地缓存
            localEntries := make([]*KVEntry, 0, len(ks))
            remainingKeys := make([]string, 0, len(ks))
            
            for _, key := range ks {
                if val, ok := c.localCache.Load(key); ok {
                    localEntries = append(localEntries, val.(*KVEntry))
                } else {
                    remainingKeys = append(remainingKeys, key)
                }
            }
            
            // 批量从Redis加载
            if len(remainingKeys) > 0 {
                pipe := c.redisClients[idx].Pipeline()
                cmds := make([]*redis.StringCmd, len(remainingKeys))
                for i, key := range remainingKeys {
                    cmds[i] = pipe.Get(ctx, key)
                }
                _, err := pipe.Exec(ctx)
                if err == nil {
                    for i, cmd := range cmds {
                        data, err := cmd.Bytes()
                        if err == nil {
                            entry := deserializeEntry(data)
                            localEntries = append(localEntries, entry)
                        }
                    }
                }
            }
            
            mu.Lock()
            results = append(results, localEntries...)
            mu.Unlock()
        }(nodeIdx, nodeKeys)
    }
    
    wg.Wait()
    return results, nil
}

// 辅助函数：生成缓存键
func generateCacheKey(layerID, headID, position int) string {
    buf := make([]byte, 12)
    binary.BigEndian.PutUint32(buf[0:4], uint32(layerID))
    binary.BigEndian.PutUint32(buf[4:8], uint32(headID))
    binary.BigEndian.PutUint32(buf[8:12], uint32(position))
    return string(buf)
}

package ringattention

import (
    "context"
    "log"
    "sync"
    "time"
    
    "github.com/yourorg/distributed-kv/kvstore"
)

// RingAttentionEngine 环形注意力计算引擎
type RingAttentionEngine struct {
    nodeID       int           // 当前节点ID
    totalNodes   int           // 总节点数
    kvCache      *kvstore.DistributedKVCache
    computeFunc  AttentionCompute // 实际的注意力计算函数
    config       RingConfig
}

type RingConfig struct {
    ChunkSize    int // 每个chunk的token数
    OverlapSize  int // 重叠区域大小，用于边界处理
    Timeout      time.Duration
    MaxRetries   int
}

// AttentionCompute 注意力计算函数签名
type AttentionCompute func(q, k, v [][]float16) [][]float16

// RingState 环形通信状态
type RingState struct {
    currentChunk int       // 当前处理的chunk编号
    kvBuffer     []*kvstore.KVEntry // 接收到的KV数据缓冲区
    resultBuffer [][]float16 // 局部注意力结果
    mu           sync.Mutex
}

// ExecuteRingAttention 执行完整的环形注意力计算
func (e *RingAttentionEngine) ExecuteRingAttention(ctx context.Context, 
    query [][]float16) ([][]float16, error) {
    
    // 1. 计算总chunk数
    totalChunks := len(query) / e.config.ChunkSize
    if len(query)%e.config.ChunkSize != 0 {
        totalChunks++
    }
    
    // 2. 初始化环形状态
    state := &RingState{
        currentChunk: e.nodeID, // 从当前节点负责的chunk开始
        kvBuffer:     make([]*kvstore.KVEntry, 0, e.config.ChunkSize),
        resultBuffer: make([][]float16, 0),
    }
    
    // 3. 创建通信管道
    sendCh := make(chan []*kvstore.KVEntry, e.totalNodes)
    recvCh := make(chan []*kvstore.KVEntry, e.totalNodes)
    
    // 4. 启动发送goroutine
    go e.sendLoop(ctx, sendCh)
    
    // 5. 启动接收goroutine
    go e.recvLoop(ctx, recvCh)
    
    // 6. 主循环：处理所有chunk
    for step := 0; step < totalChunks; step++ {
        select {
        case <-ctx.Done():
            return nil, ctx.Err()
        default:
        }
        
        // 6.1 加载当前chunk的KV数据
        chunkKV, err := e.loadChunkKV(ctx, state.currentChunk)
        if err != nil {
            log.Printf("Failed to load chunk %d: %v", state.currentChunk, err)
            return nil, err
        }
        
        // 6.2 发送当前chunk到下一个节点
        sendCh <- chunkKV
        
        // 6.3 接收上一个节点的chunk数据
        receivedKV := <-recvCh
        
        // 6.4 计算当前chunk的注意力
        chunkResult := e.computeChunkAttention(
            query[state.currentChunk*e.config.ChunkSize:min((state.currentChunk+1)*e.config.ChunkSize, len(query))],
            receivedKV,
        )
        
        // 6.5 累加结果
        state.mu.Lock()
        state.resultBuffer = append(state.resultBuffer, chunkResult...)
        state.mu.Unlock()
        
        // 6.6 移动到下一个chunk（环形）
        state.currentChunk = (state.currentChunk + 1) % totalChunks
    }
    
    // 7. 合并所有节点的结果
    finalResult := e.mergeResults(state.resultBuffer, totalChunks)
    
    return finalResult, nil
}

// computeChunkAttention 计算单个chunk的注意力
func (e *RingAttentionEngine) computeChunkAttention(
    queryChunk [][]float16,
    kvChunk []*kvstore.KVEntry,
) [][]float16 {
    
    // 提取键和值
    keys := make([][]float16, len(kvChunk))
    values := make([][]float16, len(kvChunk))
    for i, entry := range kvChunk {
        keys[i] = entry.KeyData
        values[i] = entry.ValueData
    }
    
    // 调用实际注意力计算
    return e.computeFunc(queryChunk, keys, values)
}

// sendLoop 发送循环：将数据发送到下一个节点
func (e *RingAttentionEngine) sendLoop(ctx context.Context, sendCh <-chan []*kvstore.KVEntry) {
    nextNode := (e.nodeID + 1) % e.totalNodes
    
    for {
        select {
        case <-ctx.Done():
            return
        case data := <-sendCh:
            // 序列化并通过网络发送到下一个节点
            err := e.sendToNode(ctx, nextNode, data)
            if err != nil {
                log.Printf("Failed to send to node %d: %v", nextNode, err)
                // 重试逻辑
                for retry := 0; retry < e.config.MaxRetries; retry++ {
                    err = e.sendToNode(ctx, nextNode, data)
                    if err == nil {
                        break
                    }
                    time.Sleep(time.Millisecond * 100 * (1 << retry))
                }
            }
        }
    }
}

// recvLoop 接收循环：从上一个节点接收数据
func (e *RingAttentionEngine) recvLoop(ctx context.Context, recvCh chan<- []*kvstore.KVEntry) {
    prevNode := (e.nodeID - 1 + e.totalNodes) % e.totalNodes
    
    for {
        select {
        case <-ctx.Done():
            return
        default:
            data, err := e.receiveFromNode(ctx, prevNode)
            if err == nil {
                recvCh <- data
            }
            // 短暂休眠避免忙等待
            time.Sleep(time.Microsecond * 10)
        }
    }
}

package scheduler

import (
    "sync"
    "time"
)

// AttentionMode 注意力模式枚举
type AttentionMode int

const (
    ModeFull        AttentionMode = iota // 全局注意力
    ModeSlidingWindow                    // 滑动窗口注意力
    ModeSparse                           // 稀疏注意力
    ModeHybrid                           // 混合模式
)

// SparseAttentionScheduler 稀疏注意力调度器
type SparseAttentionScheduler struct {
    config           SchedulerConfig
    attentionHistory map[int]float64 // 记录每个位置的注意力权重历史
    mu               sync.RWMutex
}

type SchedulerConfig struct {
    WindowSize       int     // 滑动窗口大小
    GlobalTokenRatio float64 // 全局token比例（0-1）
    SparsityFactor   int     // 稀疏因子：每隔几个token保留一个
    HistoryLength    int     // 历史记录长度
}

// DecideAttentionPattern 根据位置和历史决定注意力模式
func (s *SparseAttentionScheduler) DecideAttentionPattern(
    position int,
    totalLength int,
    queryEmbedding []float16,
) AttentionMode {
    
    // 1. 特殊位置使用全局注意力
    if position == 0 || position == totalLength-1 {
        return ModeFull
    }
    
    // 2. 根据位置重要性决定
    importance := s.calculateImportance(position)
    
    if importance > 0.8 {
        // 高重要性位置使用全局注意力
        return ModeFull
    } else if importance > 0.5 {
        // 中等重要性使用混合模式
        return ModeHybrid
    } else {
        // 低重要性使用稀疏模式
        return ModeSparse
    }
}

// calculateImportance 计算某个位置的重要性分数
func (s *SparseAttentionScheduler) calculateImportance(position int) float64 {
    s.mu.RLock()
    defer s.mu.RUnlock()
    
    // 基于历史注意力权重计算
    totalWeight := 0.0
    count := 0
    
    for pos, weight := range s.attentionHistory {
        if abs(pos-position) <= s.config.WindowSize {
            totalWeight += weight
            count++
        }
    }
    
    if count == 0 {
        return 0.5 // 默认中等重要性
    }
    
    return totalWeight / float64(count)
}

// GenerateAttentionMask 根据模式生成注意力掩码
func (s *SparseAttentionScheduler) GenerateAttentionMask(
    mode AttentionMode,
    queryPos int,
    keyPositions []int,
) []bool {
    
    mask := make([]bool, len(keyPositions))
    
    switch mode {
    case ModeFull:
        // 全部可见
        for i := range mask {
            mask[i] = true
        }
        
    case ModeSlidingWindow:
        // 只关注窗口内的token
        for i, pos := range keyPositions {
            if abs(pos-queryPos) <= s.config.WindowSize {
                mask[i] = true
            }
        }
        
    case ModeSparse:
        // 每隔sparsityFactor个token保留一个
        for i, pos := range keyPositions {
            if pos%s.config.SparsityFactor == 0 || abs(pos-queryPos) <= s.config.WindowSize {
                mask[i] = true
            }
        }
        
    case ModeHybrid:
        // 全局token + 滑动窗口
        for i, pos := range keyPositions {
            if pos == 0 || pos == len(keyPositions)-1 || 
               abs(pos-queryPos) <= s.config.WindowSize {
                mask[i] = true
            }
        }
    }
    
    return mask
}

// UpdateAttentionHistory 更新注意力权重历史
func (s *SparseAttentionScheduler) UpdateAttentionHistory(
    positions []int,
    weights []float64,
) {
    s.mu.Lock()
    defer s.mu.Unlock()
    
    for i, pos := range positions {
        if i < len(weights) {
            s.attentionHistory[pos] = weights[i]
        }
    }
    
    // 清理过时历史
    if len(s.attentionHistory) > s.config.HistoryLength {
        for pos := range s.attentionHistory {
            if len(s.attentionHistory) <= s.config.HistoryLength {
                break
            }
            delete(s.attentionHistory, pos)
        }
    }
}