// internal/types/frame.go
package types

import (
    "image"
    "time"
)

// VideoFrame 表示一帧视频数据
type VideoFrame struct {
    // 帧序号，用于追踪和调试
    Seq uint64
    // 显示时间戳（Presentation Timestamp），单位微秒
    PTS int64
    // 采集时间，用于延迟计算
    CaptureTime time.Time
    // 原始图像数据，可能为nil（如果解码失败）
    Image *image.RGBA
    // 解码后的张量数据，形状为 [C, H, W]
    Tensor []float32
    // 元数据，用于扩展
    Meta map[string]interface{}
}

// internal/types/result.go
package types

// VideoResult 表示视频理解的结果
type VideoResult struct {
    // 结果ID，用于追踪
    ID string
    // 描述文本
    Description string
    // 开始时间戳（对应视频中的起始帧）
    StartPTS int64
    // 结束时间戳（对应视频中的结束帧）
    EndPTS int64
    // 置信度 [0, 1]
    Confidence float32
    // 结构化数据（可选）
    StructuredData map[string]interface{}
    // 生成时间
    CreatedAt time.Time
}

// pkg/utils/ring_buffer.go
package utils

import (
    "sync"
    "video-understanding/internal/types"
)

// RingBuffer 线程安全的环形缓冲区
// 用于在采集线程和推理线程之间传递视频帧
type RingBuffer struct {
    mu     sync.Mutex
    buffer []*types.VideoFrame
    size   int
    head   int // 读指针
    tail   int // 写指针
    count  int // 当前元素数量
    notFull  *sync.Cond // 缓冲区未满条件变量
    notEmpty *sync.Cond // 缓冲区非空条件变量
}

// NewRingBuffer 创建一个新的环形缓冲区
// size: 缓冲区容量，建议设置为模型最大处理帧数的2倍
func NewRingBuffer(size int) *RingBuffer {
    rb := &RingBuffer{
        buffer: make([]*types.VideoFrame, size),
        size:   size,
    }
    rb.notFull = sync.NewCond(&rb.mu)
    rb.notEmpty = sync.NewCond(&rb.mu)
    return rb
}

// Push 向缓冲区写入一帧数据
// 如果缓冲区已满，会阻塞直到有空间
func (rb *RingBuffer) Push(frame *types.VideoFrame) {
    rb.mu.Lock()
    defer rb.mu.Unlock()

    // 缓冲区满时等待
    for rb.count == rb.size {
        rb.notFull.Wait()
    }

    // 写入数据
    rb.buffer[rb.tail] = frame
    rb.tail = (rb.tail + 1) % rb.size
    rb.count++

    // 通知等待的消费者
    rb.notEmpty.Signal()
}

// Pop 从缓冲区读取一帧数据
// 如果缓冲区为空，会阻塞直到有新数据
func (rb *RingBuffer) Pop() *types.VideoFrame {
    rb.mu.Lock()
    defer rb.mu.Unlock()

    // 缓冲区空时等待
    for rb.count == 0 {
        rb.notEmpty.Wait()
    }

    // 读取数据
    frame := rb.buffer[rb.head]
    rb.buffer[rb.head] = nil // 防止内存泄漏
    rb.head = (rb.head + 1) % rb.size
    rb.count--

    // 通知等待的生产者
    rb.notFull.Signal()

    return frame
}

// TryPush 非阻塞写入，成功返回true，失败返回false
func (rb *RingBuffer) TryPush(frame *types.VideoFrame) bool {
    rb.mu.Lock()
    defer rb.mu.Unlock()

    if rb.count == rb.size {
        return false
    }

    rb.buffer[rb.tail] = frame
    rb.tail = (rb.tail + 1) % rb.size
    rb.count++
    rb.notEmpty.Signal()
    return true
}

// Len 返回当前缓冲区中的数据量
func (rb *RingBuffer) Len() int {
    rb.mu.Lock()
    defer rb.mu.Unlock()
    return rb.count
}

// internal/inference/kv_cache.go
package inference

import (
    "sync"
    "video-understanding/internal/types"
)

// KVCache 管理推理过程中使用的Key-Value缓存
// 基于滑动窗口机制，只保留最近N帧的KV向量
type KVCache struct {
    mu         sync.RWMutex
    maxFrames  int           // 最大缓存帧数，对应滑动窗口大小
    cache      map[uint64]*FrameKV // 帧序号 -> KV向量
    frameOrder []uint64      // 帧序号顺序列表，用于LRU淘汰
    currentSize int          // 当前缓存的帧数
}

// FrameKV 存储单帧的KV向量
type FrameKV struct {
    // Key向量，形状为 [num_layers, num_heads, seq_len, head_dim]
    Key [][][][]float32
    // Value向量，形状同上
    Value [][][][]float32
    // 该帧的时间戳
    PTS int64
}

// NewKVCache 创建新的KV缓存
// maxFrames: 最大缓存帧数，建议设置为模型滑动窗口大小
func NewKVCache(maxFrames int) *KVCache {
    return &KVCache{
        maxFrames:  maxFrames,
        cache:      make(map[uint64]*FrameKV),
        frameOrder: make([]uint64, 0, maxFrames),
    }
}

// Insert 插入新帧的KV向量
// 如果缓存已满，会淘汰最旧的帧
func (kc *KVCache) Insert(seq uint64, kv *FrameKV) {
    kc.mu.Lock()
    defer kc.mu.Unlock()

    // 如果已存在，直接更新
    if _, exists := kc.cache[seq]; exists {
        kc.cache[seq] = kv
        return
    }

    // 如果缓存已满，淘汰最旧的帧
    if kc.currentSize >= kc.maxFrames {
        oldestSeq := kc.frameOrder[0]
        delete(kc.cache, oldestSeq)
        kc.frameOrder = kc.frameOrder[1:]
        kc.currentSize--
    }

    // 插入新帧
    kc.cache[seq] = kv
    kc.frameOrder = append(kc.frameOrder, seq)
    kc.currentSize++
}

// Get 获取指定帧的KV向量
// 如果不存在，返回nil
func (kc *KVCache) Get(seq uint64) *FrameKV {
    kc.mu.RLock()
    defer kc.mu.RUnlock()

    return kc.cache[seq]
}

// GetAll 获取所有缓存的KV向量，按帧序号排序
// 返回结果用于模型推理时的Key-Value拼接
func (kc *KVCache) GetAll() []*FrameKV {
    kc.mu.RLock()
    defer kc.mu.RUnlock()

    result := make([]*FrameKV, 0, kc.currentSize)
    for _, seq := range kc.frameOrder {
        if kv, exists := kc.cache[seq]; exists {
            result = append(result, kv)
        }
    }
    return result
}

// Clear 清空缓存
func (kc *KVCache) Clear() {
    kc.mu.Lock()
    defer kc.mu.Unlock()

    kc.cache = make(map[uint64]*FrameKV)
    kc.frameOrder = make([]uint64, 0, kc.maxFrames)
    kc.currentSize = 0
}

// internal/inference/engine.go
package inference

import (
    "context"
    "fmt"
    "log"
    "sync"
    "time"
    "video-understanding/internal/types"
    "video-understanding/pkg/utils"
)

// InferenceEngine 推理引擎，负责加载模型并执行流式推理
type InferenceEngine struct {
    // 模型路径
    modelPath string
    // 推理后端（ONNX Runtime / TensorRT）
    backend string
    // 滑动窗口大小（帧数）
    windowSize int
    // 帧采样间隔（每隔多少帧处理一帧）
    frameInterval int
    // 输入队列
    inputChan chan *types.VideoFrame
    // 输出队列
    outputChan chan *types.VideoResult
    // KV缓存
    kvCache *KVCache
    // 环形缓冲区（用于采集与推理的解耦）
    ringBuffer *utils.RingBuffer
    // 上下文，用于取消
    ctx context.Context
    cancel context.CancelFunc
    // 等待组
    wg sync.WaitGroup
    // 模型句柄（实际实现中为ONNX Runtime的Session）
    modelHandle interface{}
}

// NewInferenceEngine 创建推理引擎实例
func NewInferenceEngine(modelPath string, windowSize int, frameInterval int) *InferenceEngine {
    ctx, cancel := context.WithCancel(context.Background())
    return &InferenceEngine{
        modelPath:     modelPath,
        windowSize:    windowSize,
        frameInterval: frameInterval,
        inputChan:     make(chan *types.VideoFrame, 100),
        outputChan:    make(chan *types.VideoResult, 50),
        kvCache:       NewKVCache(windowSize),
        ringBuffer:    utils.NewRingBuffer(windowSize * 2),
        ctx:           ctx,
        cancel:        cancel,
    }
}

// Start 启动推理引擎
// 会启动两个goroutine：一个用于帧处理，一个用于结果生成
func (e *InferenceEngine) Start() error {
    // 加载模型
    if err := e.loadModel(); err != nil {
        return fmt.Errorf("加载模型失败: %w", err)
    }

    log.Printf("推理引擎启动，窗口大小=%d帧，采样间隔=%d帧", e.windowSize, e.frameInterval)

    // 启动帧处理协程
    e.wg.Add(1)
    go e.frameProcessor()

    // 启动结果生成协程
    e.wg.Add(1)
    go e.resultGenerator()

    return nil
}

// Stop 停止推理引擎
func (e *InferenceEngine) Stop() {
    e.cancel()
    e.wg.Wait()
    log.Println("推理引擎已停止")
}

// SubmitFrame 提交一帧数据到推理引擎
// 外部采集模块调用此方法
func (e *InferenceEngine) SubmitFrame(frame *types.VideoFrame) {
    select {
    case e.inputChan <- frame:
    default:
        // 队列满时丢弃最旧的帧
        log.Printf("输入队列已满，丢弃帧 %d", frame.Seq)
    }
}

// GetOutput 获取推理结果（阻塞）
func (e *InferenceEngine) GetOutput() *types.VideoResult {
    select {
    case result := <-e.outputChan:
        return result
    case <-e.ctx.Done():
        return nil
    }
}

// frameProcessor 帧处理协程
// 负责从输入通道获取帧，进行预处理后加入KV缓存
func (e *InferenceEngine) frameProcessor() {
    defer e.wg.Done()

    frameCount := 0
    for {
        select {
        case <-e.ctx.Done():
            return
        case frame := <-e.inputChan:
            // 帧采样控制：每隔frameInterval帧处理一帧
            frameCount++
            if frameCount%e.frameInterval != 0 {
                continue
            }

            // 预处理：图像归一化、转换为张量等
            preprocessed := e.preprocess(frame)

            // 执行模型前向传播（实际调用ONNX Runtime）
            kv := e.forward(preprocessed)

            // 更新KV缓存
            e.kvCache.Insert(frame.Seq, kv)

            // 当缓存达到窗口大小时，触发推理
            if e.kvCache.currentSize >= e.windowSize {
                e.triggerInference()
            }
        }
    }
}

// resultGenerator 结果生成协程
// 负责从模型输出中解析描述文本，并推送结果
func (e *InferenceEngine) resultGenerator() {
    defer e.wg.Done()

    for {
        select {
        case <-e.ctx.Done():
            return
        default:
            // 从模型获取输出（非阻塞）
            if e.hasModelOutput() {
                result := e.parseModelOutput()
                select {
                case e.outputChan <- result:
                default:
                    log.Printf("输出队列已满，丢弃结果")
                }
            }
            // 控制轮询频率，避免CPU空转
            time.Sleep(10 * time.Millisecond)
        }
    }
}

// loadModel 加载模型（伪代码，实际实现需调用推理后端API）
func (e *InferenceEngine) loadModel() error {
    log.Printf("正在加载模型: %s", e.modelPath)
    // TODO: 调用ONNX Runtime或TensorRT加载模型
    // session, err := ort.NewSession(e.modelPath, ort.SessionOptions{})
    // if err != nil { return err }
    // e.modelHandle = session
    return nil
}

// preprocess 帧预处理
func (e *InferenceEngine) preprocess(frame *types.VideoFrame) *types.VideoFrame {
    // TODO: 执行图像缩放、归一化、转换为张量
    // 此处简化处理，直接返回原帧
    return frame
}

// forward 执行模型前向传播（伪代码）
func (e *InferenceEngine) forward(frame *types.VideoFrame) *FrameKV {
    // TODO: 调用推理后端执行前向传播
    // 返回该帧的KV向量
    return &FrameKV{}
}

// triggerInference 触发推理
// 当累积足够帧后，执行完整的推理过程
func (e *InferenceEngine) triggerInference() {
    // 获取所有缓存的KV向量
    allKV := e.kvCache.GetAll()
    
    // TODO: 将KV向量拼接后输入模型解码器
    // 执行自回归解码，生成描述文本
    log.Printf("触发推理，当前缓存帧数: %d", len(allKV))
}

// hasModelOutput 检查模型是否有输出
func (e *InferenceEngine) hasModelOutput() bool {
    // TODO: 检查模型解码器是否已经生成完整描述
    return false
}

// parseModelOutput 解析模型输出
func (e *InferenceEngine) parseModelOutput() *types.VideoResult {
    // TODO: 从模型输出中解析描述文本
    return &types.VideoResult{
        ID:          fmt.Sprintf("result-%d", time.Now().UnixNano()),
        Description: "检测到一个人正在跑步",
        StartPTS:    0,
        EndPTS:      1000,
        Confidence:  0.95,
        CreatedAt:   time.Now(),
    }
}

// cmd/server.go
package main

import (
    "encoding/json"
    "log"
    "net/http"
    "time"
    "video-understanding/internal/inference"
    "video-understanding/internal/types"
)

func main() {
    // 配置参数
    config := &types.Config{
        ModelPath:     "/models/gpt4v.onnx",
        WindowSize:    32,   // 滑动窗口32帧
        FrameInterval: 3,    // 每隔3帧处理一帧（10fps）
        Port:          ":8080",
    }

    // 创建推理引擎
    engine := inference.NewInferenceEngine(
        config.ModelPath,
        config.WindowSize,
        config.FrameInterval,
    )

    // 启动引擎
    if err := engine.Start(); err != nil {
        log.Fatalf("启动推理引擎失败: %v", err)
    }
    defer engine.Stop()

    // 启动结果消费协程
    go func() {
        for {
            result := engine.GetOutput()
            if result == nil {
                continue
            }
            // 将结果写入日志或推送至消息队列
            log.Printf("推理结果: %+v", result)
        }
    }()

    // 启动HTTP服务
    http.HandleFunc("/api/v1/upload", func(w http.ResponseWriter, r *http.Request) {
        // 处理视频文件上传
        w.WriteHeader(http.StatusOK)
        json.NewEncoder(w).Encode(map[string]string{
            "message": "视频上传成功",
        })
    })

    http.HandleFunc("/api/v1/stream", func(w http.ResponseWriter, r *http.Request) {
        // 处理实时视频流
        // 使用WebSocket或SSE推送结果
        w.WriteHeader(http.StatusOK)
        json.NewEncoder(w).Encode(map[string]string{
            "message": "实时流处理中",
        })
    })

    log.Printf("服务启动在端口 %s", config.Port)
    log.Fatal(http.ListenAndServe(config.Port, nil))
}

// 伪代码：加载量化模型
options := ort.NewSessionOptions()
options.SetIntraOpNumThreads(4)
options.SetGraphOptimizationLevel(ort.GraphOptimizationLevelAll)

// 加载FP16模型
session, err := ort.NewSession("model_fp16.onnx", options)

// 或加载INT8模型
session, err := ort.NewSession("model_int8.onnx", options)

// 启用所有图优化
options.SetGraphOptimizationLevel(ort.GraphOptimizationLevelAll)

// 启用特定的优化pass
options.AddConfigEntry("optimization.level", "99")
options.AddConfigEntry("enable_pattern_optimization", "1")

// dynamic_batch.go
type DynamicBatcher struct {
    mu          sync.Mutex
    pending     []*types.VideoFrame
    batchSize   int
    maxWaitTime time.Duration
}

func (b *DynamicBatcher) Add(frame *types.VideoFrame) {
    b.mu.Lock()
    defer b.mu.Unlock()
    
    b.pending = append(b.pending, frame)
    
    // 达到batch大小或等待超时时触发推理
    if len(b.pending) >= b.batchSize {
        b.flush()
    }
}

func (b *DynamicBatcher) flush() {
    // 将pending中的帧组成batch，执行推理
    // 注意：不同视频的帧需要保持独立的时间轴
    batch := make([]*types.VideoFrame, len(b.pending))
    copy(batch, b.pending)
    b.pending = b.pending[:0]
    
    go b.executeBatch(batch)
}

// tensor_pool.go
type TensorPool struct {
    pool sync.Pool
}

func NewTensorPool(shape []int64) *TensorPool {
    return &TensorPool{
        pool: sync.Pool{
            New: func() interface{} {
                // 创建固定大小的张量
                size := int64(1)
                for _, dim := range shape {
                    size *= dim
                }
                return make([]float32, size)
            },
        },
    }
}

func (p *TensorPool) Get() []float32 {
    return p.pool.Get().([]float32)
}

func (p *TensorPool) Put(tensor []float32) {
    // 清零后放回池中
    for i := range tensor {
        tensor[i] = 0
    }
    p.pool.Put(tensor)
}

// 使用unsafe.Pointer传递张量数据
// 注意：确保在数据使用期间不会被GC回收
func preprocess(frame *types.VideoFrame) *float32 {
    // 直接在原始图像数据上进行归一化
    // 返回指向处理后的张量的指针
    return &frame.Tensor[0]
}

// 推理时直接使用该指针
session.Run(ort.NewValue(inputTensorPtr, inputShape), ...)

// worker_pool.go
type WorkerPool struct {
    workers  int
    taskChan chan func()
    wg       sync.WaitGroup
}

func NewWorkerPool(workers int) *WorkerPool {
    pool := &WorkerPool{
        workers:  workers,
        taskChan: make(chan func(), 100),
    }
    
    for i := 0; i < workers; i++ {
        pool.wg.Add(1)
        go pool.worker()
    }
    return pool
}

func (p *WorkerPool) Submit(task func()) {
    p.taskChan <- task
}

func (p *WorkerPool) worker() {
    defer p.wg.Done()
    for task := range p.taskChan {
        task()
    }
}