"""
Arrow-based buffer management for zero-copy operations
"""
from typing import Dict, Optional, Any
import pyarrow as pa
import numpy as np

class StreamBuffer:
    """Manages Arrow buffers for zero-copy operations"""
    
    def __init__(self):
        self.active_buffers: Dict[str, pa.Buffer] = {}
        self.buffer_refs: Dict[str, Any] = {}  # Keep references to prevent GC
        
    def register_buffer(self, 
                       name: str, 
                       arrow_table: pa.Table,
                       pin: bool = False) -> str:
        """Register an Arrow buffer for zero-copy access"""
        buffer_id = f"{name}_{len(self.active_buffers)}"
        
        # Get the buffer from Arrow table
        # Assuming first column contains our data
        self.active_buffers[buffer_id] = arrow_table.column(0).chunks[0].buffers()[1]
        
        if pin:
            # Keep a reference to prevent garbage collection
            self.buffer_refs[buffer_id] = arrow_table
            
        return buffer_id
        
    def get_buffer(self, buffer_id: str) -> Optional[pa.Buffer]:
        """Get a registered buffer"""
        return self.active_buffers.get(buffer_id)
        
    def get_numpy_view(self, 
                      buffer_id: str,
                      dtype: np.dtype = np.float32) -> np.ndarray:
        """Get a numpy view of the buffer without copying"""
        buffer = self.get_buffer(buffer_id)
        if buffer is None:
            raise KeyError(f"Buffer {buffer_id} not found")
            
        # Create numpy view without copying
        return np.frombuffer(buffer, dtype=dtype)
        
    def release_buffer(self, buffer_id: str):
        """Release a buffer"""
        self.active_buffers.pop(buffer_id, None)
        self.buffer_refs.pop(buffer_id, None)
        
    def clear_buffers(self):
        """Clear all buffers"""
        self.active_buffers.clear()
        self.buffer_refs.clear()
