// internal/types/frame.go
package types

import (
    "image"
    "time"
)

// VideoFrame represents a single video frame
type VideoFrame struct {
    // Frame sequence number for tracking and debugging
    Seq uint64
    // Presentation Timestamp in microseconds
    PTS int64
    // Capture time for latency calculation
    CaptureTime time.Time
    // Original image data, may be nil (if decoding failed)
    Image *image.RGBA
    // Decoded tensor data with shape [C, H, W]
    Tensor []float32
    // Metadata for extension
    Meta map[string]interface{}
}

// internal/types/result.go
package types

// VideoResult represents the result of video understanding
type VideoResult struct {
    // Result ID for tracking
    ID string
    // Description text
    Description string
    // Start timestamp (corresponding to the starting frame in the video)
    StartPTS int64
    // End timestamp (corresponding to the ending frame in the video)
    EndPTS int64
    // Confidence [0, 1]
    Confidence float32
    // Structured data (optional)
    StructuredData map[string]interface{}
    // Generation time
    CreatedAt time.Time
}

// pkg/utils/ring_buffer.go
package utils

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

// RingBuffer is a thread-safe ring buffer
// Used to pass video frames between capture and inference threads
type RingBuffer struct {
    mu     sync.Mutex
    buffer []*types.VideoFrame
    size   int
    head   int // read pointer
    tail   int // write pointer
    count  int // current number of elements
    notFull  *sync.Cond // condition variable for buffer not full
    notEmpty *sync.Cond // condition variable for buffer not empty
}

// NewRingBuffer creates a new ring buffer
// size: buffer capacity, recommended to be 2x the model's maximum processing frame count
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 writes a frame to the buffer
// Blocks if the buffer is full until space becomes available
func (rb *RingBuffer) Push(frame *types.VideoFrame) {
    rb.mu.Lock()
    defer rb.mu.Unlock()

    // Wait if buffer is full
    for rb.count == rb.size {
        rb.notFull.Wait()
    }

    // Write data
    rb.buffer[rb.tail] = frame
    rb.tail = (rb.tail + 1) % rb.size
    rb.count++

    // Signal waiting consumers
    rb.notEmpty.Signal()
}

// Pop reads a frame from the buffer
// Blocks if the buffer is empty until new data arrives
func (rb *RingBuffer) Pop() *types.VideoFrame {
    rb.mu.Lock()
    defer rb.mu.Unlock()

    // Wait if buffer is empty
    for rb.count == 0 {
        rb.notEmpty.Wait()
    }

    // Read data
    frame := rb.buffer[rb.head]
    rb.buffer[rb.head] = nil // prevent memory leak
    rb.head = (rb.head + 1) % rb.size
    rb.count--

    // Signal waiting producers
    rb.notFull.Signal()

    return frame
}

// TryPush performs a non-blocking write, returns true on success, false on failure
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 returns the current number of elements in the buffer
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 manages Key-Value caches used during inference
// Based on a sliding window mechanism, only retaining KV vectors from the most recent N frames
type KVCache struct {
    mu         sync.RWMutex
    maxFrames  int           // maximum cached frames, corresponds to sliding window size
    cache      map[uint64]*FrameKV // frame sequence number -> KV vector
    frameOrder []uint64      // ordered list of frame sequence numbers for LRU eviction
    currentSize int          // current number of cached frames
}

// FrameKV stores KV vectors for a single frame
type FrameKV struct {
    // Key vectors, shape [num_layers, num_heads, seq_len, head_dim]
    Key [][][][]float32
    // Value vectors, same shape
    Value [][][][]float32
    // Timestamp of this frame
    PTS int64
}

// NewKVCache creates a new KV cache
// maxFrames: maximum cached frames, recommended to be set to the model's sliding window size
func NewKVCache(maxFrames int) *KVCache {
    return &KVCache{
        maxFrames:  maxFrames,
        cache:      make(map[uint64]*FrameKV),
        frameOrder: make([]uint64, 0, maxFrames),
    }
}

// Insert inserts KV vectors for a new frame
// If the cache is full, the oldest frame is evicted
func (kc *KVCache) Insert(seq uint64, kv *FrameKV) {
    kc.mu.Lock()
    defer kc.mu.Unlock()

    // If already exists, update directly
    if _, exists := kc.cache[seq]; exists {
        kc.cache[seq] = kv
        return
    }

    // If cache is full, evict the oldest frame
    if kc.currentSize >= kc.maxFrames {
        oldestSeq := kc.frameOrder[0]
        delete(kc.cache, oldestSeq)
        kc.frameOrder = kc.frameOrder[1:]
        kc.currentSize--
    }

    // Insert new frame
    kc.cache[seq] = kv
    kc.frameOrder = append(kc.frameOrder, seq)
    kc.currentSize++
}

// Get retrieves KV vectors for a specified frame
// Returns nil if not found
func (kc *KVCache) Get(seq uint64) *FrameKV {
    kc.mu.RLock()
    defer kc.mu.RUnlock()

    return kc.cache[seq]
}

// GetAll retrieves all cached KV vectors, sorted by frame sequence number
// The returned result is used for Key-Value concatenation during model inference
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 empties the cache
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 is the inference engine, responsible for loading the model and executing streaming inference
type InferenceEngine struct {
    // Model path
    modelPath string
    // Inference backend (ONNX Runtime / TensorRT)
    backend string
    // Sliding window size (frames)
    windowSize int
    // Frame sampling interval (process one frame every N frames)
    frameInterval int
    // Input channel
    inputChan chan *types.VideoFrame
    // Output channel
    outputChan chan *types.VideoResult
    // KV cache
    kvCache *KVCache
    // Ring buffer (for decoupling capture and inference)
    ringBuffer *utils.RingBuffer
    // Context for cancellation
    ctx context.Context
    cancel context.CancelFunc
    // Wait group
    wg sync.WaitGroup
    // Model handle (in practice, this would be an ONNX Runtime Session)
    modelHandle interface{}
}

// NewInferenceEngine creates a new inference engine instance
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 starts the inference engine
// Launches two goroutines: one for frame processing, one for result generation
func (e *InferenceEngine) Start() error {
    // Load model
    if err := e.loadModel(); err != nil {
        return fmt.Errorf("failed to load model: %w", err)
    }

    log.Printf("Inference engine started, window size=%d frames, sampling interval=%d frames", e.windowSize, e.frameInterval)

    // Start frame processing goroutine
    e.wg.Add(1)
    go e.frameProcessor()

    // Start result generation goroutine
    e.wg.Add(1)
    go e.resultGenerator()

    return nil
}

// Stop stops the inference engine
func (e *InferenceEngine) Stop() {
    e.cancel()
    e.wg.Wait()
    log.Println("Inference engine stopped")
}

// SubmitFrame submits a frame to the inference engine
// Called by the external capture module
func (e *InferenceEngine) SubmitFrame(frame *types.VideoFrame) {
    select {
    case e.inputChan <- frame:
    default:
        // Drop the oldest frame when the queue is full
        log.Printf("Input queue full, dropping frame %d", frame.Seq)
    }
}

// GetOutput retrieves an inference result (blocking)
func (e *InferenceEngine) GetOutput() *types.VideoResult {
    select {
    case result := <-e.outputChan:
        return result
    case <-e.ctx.Done():
        return nil
    }
}

// frameProcessor is the frame processing goroutine
// Responsible for getting frames from the input channel, preprocessing, and inserting into KV cache
func (e *InferenceEngine) frameProcessor() {
    defer e.wg.Done()

    frameCount := 0
    for {
        select {
        case <-e.ctx.Done():
            return
        case frame := <-e.inputChan:
            // Frame sampling control: process one frame every frameInterval frames
            frameCount++
            if frameCount%e.frameInterval != 0 {
                continue
            }

            // Preprocessing: image normalization, conversion to tensor, etc.
            preprocessed := e.preprocess(frame)

            // Execute model forward pass (actual call to ONNX Runtime)
            kv := e.forward(preprocessed)

            // Update KV cache
            e.kvCache.Insert(frame.Seq, kv)

            // When the cache reaches the window size, trigger inference
            if e.kvCache.currentSize >= e.windowSize {
                e.triggerInference()
            }
        }
    }
}

// resultGenerator is the result generation goroutine
// Responsible for parsing description text from model output and pushing results
func (e *InferenceEngine) resultGenerator() {
    defer e.wg.Done()

    for {
        select {
        case <-e.ctx.Done():
            return
        default:
            // Get model output (non-blocking)
            if e.hasModelOutput() {
                result := e.parseModelOutput()
                select {
                case e.outputChan <- result:
                default:
                    log.Printf("Output queue full, dropping result")
                }
            }
            // Control polling frequency to avoid CPU spinning
            time.Sleep(10 * time.Millisecond)
        }
    }
}

// loadModel loads the model (pseudo-code, actual implementation requires calling inference backend API)
func (e *InferenceEngine) loadModel() error {
    log.Printf("Loading model: %s", e.modelPath)
    // TODO: Call ONNX Runtime or TensorRT to load the model
    // session, err := ort.NewSession(e.modelPath, ort.SessionOptions{})
    // if err != nil { return err }
    // e.modelHandle = session
    return nil
}

// preprocess performs frame preprocessing
func (e *InferenceEngine) preprocess(frame *types.VideoFrame) *types.VideoFrame {
    // TODO: Perform image scaling, normalization, conversion to tensor
    // Simplified here, returning the original frame directly
    return frame
}

// forward executes model forward pass (pseudo-code)
func (e *InferenceEngine) forward(frame *types.VideoFrame) *FrameKV {
    // TODO: Call inference backend to execute forward pass
    // Return KV vectors for this frame
    return &FrameKV{}
}

// triggerInference triggers inference
// When enough frames have accumulated, execute the complete inference process
func (e *InferenceEngine) triggerInference() {
    // Get all cached KV vectors
    allKV := e.kvCache.GetAll()
    
    // TODO: Concatenate KV vectors and feed into model decoder
    // Execute autoregressive decoding to generate description text
    log.Printf("Triggering inference, current cached frames: %d", len(allKV))
}

// hasModelOutput checks if the model has output
func (e *InferenceEngine) hasModelOutput() bool {
    // TODO: Check if the model decoder has generated a complete description
    return false
}

// parseModelOutput parses model output
func (e *InferenceEngine) parseModelOutput() *types.VideoResult {
    // TODO: Parse description text from model output
    return &types.VideoResult{
        ID:          fmt.Sprintf("result-%d", time.Now().UnixNano()),
        Description: "Detected a person running",
        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() {
    // Configuration parameters
    config := &types.Config{
        ModelPath:     "/models/gpt4v.onnx",
        WindowSize:    32,   // Sliding window of 32 frames
        FrameInterval: 3,    // Process one frame every 3 frames (10fps)
        Port:          ":8080",
    }

    // Create inference engine
    engine := inference.NewInferenceEngine(
        config.ModelPath,
        config.WindowSize,
        config.FrameInterval,
    )

    // Start engine
    if err := engine.Start(); err != nil {
        log.Fatalf("Failed to start inference engine: %v", err)
    }
    defer engine.Stop()

    // Start result consumption goroutine
    go func() {
        for {
            result := engine.GetOutput()
            if result == nil {
                continue
            }
            // Write result to log or push to message queue
            log.Printf("Inference result: %+v", result)
        }
    }()

    // Start HTTP service
    http.HandleFunc("/api/v1/upload", func(w http.ResponseWriter, r *http.Request) {
        // Handle video file upload
        w.WriteHeader(http.StatusOK)
        json.NewEncoder(w).Encode(map[string]string{
            "message": "Video uploaded successfully",
        })
    })

    http.HandleFunc("/api/v1/stream", func(w http.ResponseWriter, r *http.Request) {
        // Handle real-time video stream
        // Use WebSocket or SSE to push results
        w.WriteHeader(http.StatusOK)
        json.NewEncoder(w).Encode(map[string]string{
            "message": "Real-time stream processing",
        })
    })

    log.Printf("Service started on port %s", config.Port)
    log.Fatal(http.ListenAndServe(config.Port, nil))
}

// Pseudo-code: Loading quantized model
options := ort.NewSessionOptions()
options.SetIntraOpNumThreads(4)
options.SetGraphOptimizationLevel(ort.GraphOptimizationLevelAll)

// Load FP16 model
session, err := ort.NewSession("model_fp16.onnx", options)

// Or load INT8 model
session, err := ort.NewSession("model_int8.onnx", options)

// Enable all graph optimizations
options.SetGraphOptimizationLevel(ort.GraphOptimizationLevelAll)

// Enable specific optimization passes
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)
    
    // Trigger inference when batch size is reached or wait timeout occurs
    if len(b.pending) >= b.batchSize {
        b.flush()
    }
}

func (b *DynamicBatcher) flush() {
    // Form a batch from pending frames and execute inference
    // Note: Frames from different videos must maintain independent timelines
    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{} {
                // Create fixed-size tensor
                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) {
    // Zero out before returning to pool
    for i := range tensor {
        tensor[i] = 0
    }
    p.pool.Put(tensor)
}

// Use unsafe.Pointer to pass tensor data
// Note: Ensure the data is not garbage collected while in use
func preprocess(frame *types.VideoFrame) *float32 {
    // Perform normalization directly on the original image data
    // Return a pointer to the processed tensor
    return &frame.Tensor[0]
}

// Use the pointer directly during inference
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