/memvid/retriever.py

"""
MemvidRetriever - Fast semantic search, QR frame extraction, and context assembly
"""

import json
import logging
from pathlib import Path
from typing import List, Dict, Any, Tuple, Optional
from concurrent.futures import ThreadPoolExecutor, as_completed
import time
from functools import lru_cache
import cv2

from .utils import (
    extract_frame, decode_qr, batch_extract_and_decode,
    extract_and_decode_cached
)
from .index import IndexManager
from .config import get_default_config

logger = logging.getLogger(__name__)


class MemvidRetriever:
    """Fast retrieval from QR code videos using semantic search"""
    
    def __init__(self, video_file: str, index_file: str, 
                 config: Optional[Dict[str, Any]] = None):
        """
        Initialize MemvidRetriever
        
        Args:
            video_file: Path to QR code video
            index_file: Path to index file
            config: Optional configuration
        """
        self.video_file = str(Path(video_file).absolute())
        self.index_file = str(Path(index_file).absolute())
        self.config = config or get_default_config()
        
        # Load index
        self.index_manager = IndexManager(self.config)
        self.index_manager.load(str(Path(index_file).with_suffix('')))
        
        # Cache for decoded frames
        self._frame_cache = {}
        self._cache_size = self.config["retrieval"]["cache_size"]
        
        # Verify video file
        self._verify_video()
        
        logger.info(f"Initialized retriever with {self.index_manager.get_stats()['total_chunks']} chunks")
    
    def _verify_video(self):
        """Verify video file is accessible"""
        cap = cv2.VideoCapture(self.video_file)
        if not cap.isOpened():
            raise ValueError(f"Cannot open video file: {self.video_file}")
        
        self.total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        self.fps = cap.get(cv2.CAP_PROP_FPS)
        cap.release()
        
        logger.info(f"Video has {self.total_frames} frames at {self.fps} fps")
    
    def search(self, query: str, top_k: int = 5) -> List[str]:
        """
        Search for relevant chunks using semantic search
        
        Args:
            query: Search query
            top_k: Number of results to return
            
        Returns:
            List of relevant text chunks
        """
        start_time = time.time()
        
        # Semantic search in index
        search_results = self.index_manager.search(query, top_k)
        
        # Extract unique frame numbers
        frame_numbers = list(set(result[2]["frame"] for result in search_results))
        
        # Decode frames in parallel
        decoded_frames = self._decode_frames_parallel(frame_numbers)
        
        # Extract text from decoded data
        results = []
        for chunk_id, distance, metadata in search_results:
            frame_num = metadata["frame"]
            if frame_num in decoded_frames:
                try:
                    chunk_data = json.loads(decoded_frames[frame_num])
                    results.append(chunk_data["text"])
                except (json.JSONDecodeError, KeyError):
                    # Fallback to metadata text
                    results.append(metadata["text"])
            else:
                # Fallback to metadata text if decode failed
                results.append(metadata["text"])
        
        elapsed = time.time() - start_time
        logger.info(f"Search completed in {elapsed:.3f}s for query: '{query[:50]}...'")
        
        return results
    
    def get_chunk_by_id(self, chunk_id: int) -> Optional[str]:
        """
        Get specific chunk by ID
        
        Args:
            chunk_id: Chunk ID
            
        Returns:
            Chunk text or None
        """
        metadata = self.index_manager.get_chunk_by_id(chunk_id)
        if metadata:
            frame_num = metadata["frame"]
            decoded = self._decode_single_frame(frame_num)
            if decoded:
                try:
                    chunk_data = json.loads(decoded)
                    return chunk_data["text"]
                except (json.JSONDecodeError, KeyError):
                    pass
            return metadata["text"]
        return None
    
    def _decode_single_frame(self, frame_number: int) -> Optional[str]:
        """Decode single frame with caching"""
        # Check cache
        if frame_number in self._frame_cache:
            return self._frame_cache[frame_number]
        
        # Decode frame
        result = extract_and_decode_cached(self.video_file, frame_number)
        
        # Update cache
        if result and len(self._frame_cache) < self._cache_size:
            self._frame_cache[frame_number] = result
        
        return result
    
    def _decode_frames_parallel(self, frame_numbers: List[int]) -> Dict[int, str]:
        """
        Decode multiple frames in parallel
        
        Args:
            frame_numbers: List of frame numbers to decode
            
        Returns:
            Dict mapping frame number to decoded data
        """
        # Check cache first
        results = {}
        uncached_frames = []
        
        for frame_num in frame_numbers:
            if frame_num in self._frame_cache:
                results[frame_num] = self._frame_cache[frame_num]
            else:
                uncached_frames.append(frame_num)
        
        if not uncached_frames:
            return results
        
        # Decode uncached frames in parallel
        max_workers = self.config["retrieval"]["max_workers"]
        decoded = batch_extract_and_decode(
            self.video_file, 
            uncached_frames, 
            max_workers=max_workers
        )
        
        # Update results and cache
        for frame_num, data in decoded.items():
            results[frame_num] = data
            if len(self._frame_cache) < self._cache_size:
                self._frame_cache[frame_num] = data
        
        return results
    
    def search_with_metadata(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
        """
        Search with full metadata
        
        Args:
            query: Search query
            top_k: Number of results
            
        Returns:
            List of result dictionaries with text, score, and metadata
        """
        start_time = time.time()
        
        # Semantic search
        search_results = self.index_manager.search(query, top_k)
        
        # Extract frame numbers
        frame_numbers = list(set(result[2]["frame"] for result in search_results))
        
        # Decode frames
        decoded_frames = self._decode_frames_parallel(frame_numbers)
        
        # Build results with metadata
        results = []
        for chunk_id, distance, metadata in search_results:
            frame_num = metadata["frame"]
            
            # Get text from decoded frame or metadata
            if frame_num in decoded_frames:
                try:
                    chunk_data = json.loads(decoded_frames[frame_num])
                    text = chunk_data["text"]
                except (json.JSONDecodeError, KeyError):
                    text = metadata["text"]
            else:
                text = metadata["text"]
            
            results.append({
                "text": text,
                "score": 1.0 / (1.0 + distance),  # Convert distance to similarity score
                "chunk_id": chunk_id,
                "frame": frame_num,
                "metadata": metadata
            })
        
        elapsed = time.time() - start_time
        logger.info(f"Search with metadata completed in {elapsed:.3f}s")
        
        return results
    
    def get_context_window(self, chunk_id: int, window_size: int = 2) -> List[str]:
        """
        Get chunk with surrounding context
        
        Args:
            chunk_id: Central chunk ID
            window_size: Number of chunks before/after to include
            
        Returns:
            List of chunks in context window
        """
        chunks = []
        
        # Get chunks in window
        for offset in range(-window_size, window_size + 1):
            target_id = chunk_id + offset
            chunk = self.get_chunk_by_id(target_id)
            if chunk:
                chunks.append(chunk)
        
        return chunks
    
    def prefetch_frames(self, frame_numbers: List[int]):
        """
        Prefetch frames into cache for faster retrieval
        
        Args:
            frame_numbers: List of frame numbers to prefetch
        """
        # Only prefetch frames not in cache
        to_prefetch = [f for f in frame_numbers if f not in self._frame_cache]
        
        if to_prefetch:
            logger.info(f"Prefetching {len(to_prefetch)} frames...")
            decoded = self._decode_frames_parallel(to_prefetch)
            logger.info(f"Prefetched {len(decoded)} frames")
    
    def clear_cache(self):
        """Clear frame cache"""
        self._frame_cache.clear()
        extract_and_decode_cached.cache_clear()
        logger.info("Cleared frame cache")
    
    def get_stats(self) -> Dict[str, Any]:
        """Get retriever statistics"""
        return {
            "video_file": self.video_file,
            "total_frames": self.total_frames,
            "fps": self.fps,
            "cache_size": len(self._frame_cache),
            "max_cache_size": self._cache_size,
            "index_stats": self.index_manager.get_stats()
        }