app/build_index.py

"""
Build FAISS index from movie embeddings
This script should be run once to create the data files needed by the API
"""
import os
import json
import numpy as np
import faiss
from openai import OpenAI
import requests
from typing import Dict, List, Optional
import time
import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
import logging
from settings import get_settings
import pickle

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Checkpoint file paths
CHECKPOINT_DIR = "app/data/checkpoints"
MOVIE_DATA_CHECKPOINT = f"{CHECKPOINT_DIR}/movie_data.pkl"
EMBEDDINGS_CHECKPOINT = f"{CHECKPOINT_DIR}/embeddings_progress.pkl"
METADATA_CHECKPOINT = f"{CHECKPOINT_DIR}/metadata_progress.pkl"

def save_checkpoint(data, filepath: str):
    """Save checkpoint data to file"""
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, 'wb') as f:
        pickle.dump(data, f)
    logger.info(f"Checkpoint saved: {filepath}")

def load_checkpoint(filepath: str):
    """Load checkpoint data from file"""
    if os.path.exists(filepath):
        with open(filepath, 'rb') as f:
            data = pickle.load(f)
        logger.info(f"Checkpoint loaded: {filepath}")
        return data
    return None

def cleanup_checkpoints():
    """Remove checkpoint files after successful completion"""
    import shutil
    if os.path.exists(CHECKPOINT_DIR):
        shutil.rmtree(CHECKPOINT_DIR)
        logger.info("Checkpoint files cleaned up")

class TMDBClient:
    """Client for TMDB API with retry and backoff"""
    
    def __init__(self, api_key: str):
        self.api_key = api_key
        self.base_url = "https://api.themoviedb.org/3"
        self.session = requests.Session()
        
    def _make_request(self, endpoint: str, params: dict = None, max_retries: int = 3) -> Optional[dict]:
        """Make API request with retry and backoff"""
        if params is None:
            params = {}
        params['api_key'] = self.api_key
        
        url = f"{self.base_url}{endpoint}"
        
        for attempt in range(max_retries):
            try:
                response = self.session.get(url, params=params, timeout=10)
                
                if response.status_code == 200:
                    return response.json()
                elif response.status_code == 429:
                    # Rate limit - wait and retry
                    wait_time = 2 ** attempt
                    logger.warning(f"Rate limited, waiting {wait_time}s before retry...")
                    time.sleep(wait_time)
                    continue
                elif response.status_code == 404:
                    logger.warning(f"Resource not found: {url}")
                    return None
                else:
                    logger.error(f"API error {response.status_code}: {response.text}")
                    
            except requests.exceptions.RequestException as e:
                logger.error(f"Request failed (attempt {attempt + 1}): {e}")
                if attempt < max_retries - 1:
                    time.sleep(2 ** attempt)
                    
        return None
    
    def get_popular_movies(self, max_pages: int = 100) -> List[int]:
        """Get movie IDs from popular movies pagination"""
        movie_ids = []
        
        for page in range(1, max_pages + 1):
            logger.info(f"Fetching popular movies page {page}/{max_pages}")
            
            data = self._make_request("/movie/popular", {"page": page})
            if not data:
                logger.error(f"Failed to fetch page {page}")
                break
                
            # Check if we've exceeded total pages
            if page > data.get('total_pages', 0):
                logger.info(f"Reached last page ({data.get('total_pages')})")
                break
                
            # Extract movie IDs
            for movie in data.get('results', []):
                movie_ids.append(movie['id'])
                
            # Rate limiting
            time.sleep(0.25)  # 4 requests per second max
            
        logger.info(f"Collected {len(movie_ids)} movie IDs from {page} pages")
        return movie_ids
    
    def get_movie_details(self, movie_id: int) -> Optional[dict]:
        """Get detailed movie information"""
        return self._make_request(f"/movie/{movie_id}")
    
    def get_movie_credits(self, movie_id: int) -> Optional[dict]:
        """Get movie cast and crew"""
        return self._make_request(f"/movie/{movie_id}/credits")

def fetch_movie_data(tmdb_client: TMDBClient, movie_ids: List[int], max_workers: int = 5) -> Dict[int, dict]:
    """Fetch detailed data for all movies with controlled parallelization"""
    movies_data = {}
    
    def fetch_single_movie(movie_id: int) -> tuple:
        """Fetch details and credits for a single movie"""
        try:
            # Get basic details
            details = tmdb_client.get_movie_details(movie_id)
            if not details:
                return movie_id, None
                
            # Get credits
            credits = tmdb_client.get_movie_credits(movie_id)
            if credits:
                details['credits'] = credits
                
            return movie_id, details
            
        except Exception as e:
            logger.error(f"Error fetching movie {movie_id}: {e}")
            return movie_id, None
    
    # Process movies in batches with controlled parallelization
    batch_size = 50
    total_movies = len(movie_ids)
    
    for i in range(0, total_movies, batch_size):
        batch = movie_ids[i:i + batch_size]
        logger.info(f"Processing batch {i//batch_size + 1}/{(total_movies-1)//batch_size + 1} ({len(batch)} movies)")
        
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = {executor.submit(fetch_single_movie, movie_id): movie_id for movie_id in batch}
            
            for future in as_completed(futures):
                movie_id, movie_data = future.result()
                if movie_data:
                    movies_data[movie_id] = movie_data
                    
        # Sleep between batches to be respectful to API
        time.sleep(1)
        
    logger.info(f"Successfully fetched data for {len(movies_data)}/{total_movies} movies")
    return movies_data

def create_composite_text(movie_data: Dict) -> str:
    """Create composite text for embedding from movie data"""
    parts = []
    
    # Title
    if movie_data.get('title'):
        parts.append(f"Title: {movie_data['title']}")
    
    # Tagline
    if movie_data.get('tagline'):
        parts.append(f"Tagline: {movie_data['tagline']}")
    
    # Overview
    if movie_data.get('overview'):
        parts.append(f"Overview: {movie_data['overview']}")
    
    # Release date
    if movie_data.get('release_date'):
        parts.append(f"Release Date: {movie_data['release_date']}")
    
    # Original language
    if movie_data.get('original_language'):
        parts.append(f"Language: {movie_data['original_language']}")
    
    # Spoken languages
    if movie_data.get('spoken_languages'):
        languages = [lang.get('iso_639_1', '') for lang in movie_data['spoken_languages'] if lang.get('iso_639_1')]
        if languages:
            parts.append(f"Spoken Languages: {', '.join(languages)}")
    
    # Genres
    if movie_data.get('genres'):
        genres = [genre['name'] for genre in movie_data['genres']]
        parts.append(f"Genres: {', '.join(genres)}")
    
    # Production companies
    if movie_data.get('production_companies'):
        companies = [company['name'] for company in movie_data['production_companies']]
        if companies:
            parts.append(f"Production Companies: {', '.join(companies)}")
    
    # Production countries
    if movie_data.get('production_countries'):
        countries = [country['name'] for country in movie_data['production_countries']]
        if countries:
            parts.append(f"Production Countries: {', '.join(countries)}")
    
    # Budget (only if > 0)
    if movie_data.get('budget') and movie_data['budget'] > 0:
        parts.append(f"Budget: ${movie_data['budget']:,}")
    
    # Popularity
    if movie_data.get('popularity'):
        parts.append(f"Popularity: {movie_data['popularity']}")
    
    # Vote average
    if movie_data.get('vote_average'):
        parts.append(f"Vote Average: {movie_data['vote_average']}")
    
    # Vote count
    if movie_data.get('vote_count'):
        parts.append(f"Vote Count: {movie_data['vote_count']}")
    
    # Director(s)
    if movie_data.get('credits', {}).get('crew'):
        directors = [person['name'] for person in movie_data['credits']['crew'] if person['job'] == 'Director']
        if directors:
            parts.append(f"Director: {', '.join(directors)}")
    
    # Top 5 cast
    if movie_data.get('credits', {}).get('cast'):
        top_cast = [person['name'] for person in movie_data['credits']['cast'][:5]]
        if top_cast:
            parts.append(f"Cast: {', '.join(top_cast)}")
    
    return " / ".join(parts)

def get_embeddings_batch(texts: List[str], client: OpenAI, model: str = "text-embedding-3-small") -> List[List[float]]:
    """Get embeddings for a batch of texts with retry"""
    max_retries = 3
    
    for attempt in range(max_retries):
        try:
            response = client.embeddings.create(
                input=texts,
                model=model
            )
            return [item.embedding for item in response.data]
        except Exception as e:
            logger.error(f"Error getting embeddings (attempt {attempt + 1}): {e}")
            if attempt < max_retries - 1:
                time.sleep(2 ** attempt)
            else:
                raise

def build_index(max_pages: int = 10, model: str = "text-embedding-3-small", use_faiss: bool = True):
    """Main function to build the FAISS index and data files"""
    settings = get_settings()
    
    # Initialize clients
    tmdb_client = TMDBClient(settings.tmdb_api_key)
    openai_client = OpenAI(api_key=settings.openai_api_key)
    
    # Create data directory
    os.makedirs("app/data", exist_ok=True)
    
    # Check for existing movie data checkpoint
    movies_data = load_checkpoint(MOVIE_DATA_CHECKPOINT)
    
    if movies_data is not None:
        logger.info(f"🔄 Resuming from checkpoint: {len(movies_data)} movies data found")
    else:
        # Step 1: Get movie IDs
        logger.info(f"Fetching movie IDs from TMDB (max {max_pages} pages)...")
        movie_ids = tmdb_client.get_popular_movies(max_pages=max_pages)
        
        if not movie_ids:
            logger.error("❌ No movie IDs retrieved from TMDB")
            return
        
        # Step 2: Fetch detailed movie data
        logger.info(f"Fetching detailed data for {len(movie_ids)} movies...")
        movies_data = fetch_movie_data(tmdb_client, movie_ids)
        
        if not movies_data:
            logger.error("❌ No movie data retrieved")
            return
        
        # Save movie data checkpoint
        save_checkpoint(movies_data, MOVIE_DATA_CHECKPOINT)
    
    # Step 3: Create composite texts and process embeddings in batches
    logger.info("Creating embeddings...")
    embeddings = []
    id_map = {}
    movie_metadata = {}
    processed_movie_ids = set()
    
    batch_size = 20  # Process 20 texts at a time
    
    # Check for existing embedding progress
    embedding_checkpoint = load_checkpoint(EMBEDDINGS_CHECKPOINT)
    metadata_checkpoint = load_checkpoint(METADATA_CHECKPOINT)
    
    if embedding_checkpoint is not None and metadata_checkpoint is not None:
        embeddings = embedding_checkpoint['embeddings']
        id_map = embedding_checkpoint['id_map']
        processed_movie_ids = set(embedding_checkpoint['processed_movie_ids'])
        movie_metadata = metadata_checkpoint
        logger.info(f"🔄 Resuming embeddings from checkpoint: {len(embeddings)} embeddings found")
    else:
        logger.info("Starting embeddings from scratch")
    
    # Process remaining movies
    remaining_movies = {k: v for k, v in movies_data.items() if k not in processed_movie_ids}
    logger.info(f"Processing {len(remaining_movies)} remaining movies")
    
    composite_texts = []
    current_movie_ids = []
    
    for movie_id, movie_data in remaining_movies.items():
        # Create composite text
        composite_text = create_composite_text(movie_data)
        composite_texts.append(composite_text)
        current_movie_ids.append(movie_id)
        
        # Store metadata
        release_year = 0
        if movie_data.get("release_date"):
            try:
                release_year = int(movie_data["release_date"][:4])
            except (ValueError, IndexError):
                release_year = 0
                
        movie_metadata[str(movie_id)] = {
            "id": movie_id,
            "title": movie_data.get("title", ""),
            "year": release_year,
            "poster_path": movie_data.get("poster_path"),
            "release_date": movie_data.get("release_date"),
            "genres": [g["name"] for g in movie_data.get("genres", [])]
        }
        
        # Process batch when full
        if len(composite_texts) >= batch_size:
            logger.info(f"Processing embedding batch ({len(embeddings)} done, {len(composite_texts)} in batch)")
            
            try:
                batch_embeddings = get_embeddings_batch(composite_texts, openai_client, model)
                embeddings.extend(batch_embeddings)
                
                # Update ID mapping and processed set
                for i, mid in enumerate(current_movie_ids):
                    id_map[str(mid)] = len(id_map)
                    processed_movie_ids.add(mid)
                
                # Save progress checkpoints
                embedding_data = {
                    'embeddings': embeddings,
                    'id_map': id_map,
                    'processed_movie_ids': list(processed_movie_ids)
                }
                save_checkpoint(embedding_data, EMBEDDINGS_CHECKPOINT)
                save_checkpoint(movie_metadata, METADATA_CHECKPOINT)
                
                # Clear batch
                composite_texts = []
                current_movie_ids = []
                
                # Sleep between batches
                time.sleep(0.5)
                
            except Exception as e:
                logger.error(f"Failed to process batch: {e}")
                logger.info("Progress has been saved, you can restart the script to resume")
                return
    
    # Process remaining texts
    if composite_texts:
        logger.info(f"Processing final embedding batch ({len(composite_texts)} texts)")
        try:
            batch_embeddings = get_embeddings_batch(composite_texts, openai_client, model)
            embeddings.extend(batch_embeddings)
            
            for i, mid in enumerate(current_movie_ids):
                id_map[str(mid)] = len(id_map)
                processed_movie_ids.add(mid)
            
            # Save final progress
            embedding_data = {
                'embeddings': embeddings,
                'id_map': id_map,
                'processed_movie_ids': list(processed_movie_ids)
            }
            save_checkpoint(embedding_data, EMBEDDINGS_CHECKPOINT)
            save_checkpoint(movie_metadata, METADATA_CHECKPOINT)
            
        except Exception as e:
            logger.error(f"Failed to process final batch: {e}")
            logger.info("Progress has been saved, you can restart the script to resume")
            return
    
    if not embeddings:
        logger.error("❌ No embeddings generated")
        return
    
    logger.info(f"Generated {len(embeddings)} embeddings")
    
    # Step 4: Save embeddings as numpy array
    embeddings_array = np.array(embeddings, dtype=np.float32)
    np.save("app/data/movies.npy", embeddings_array)
    logger.info(f"Saved embeddings matrix: {embeddings_array.shape}")
    
    # Step 5: Build and save FAISS index
    if use_faiss:
        logger.info("Building FAISS index...")
        dimension = embeddings_array.shape[1]
        
        # Choose index type based on size
        if len(embeddings) < 10000:
            # For smaller datasets, use flat index
            index = faiss.IndexFlatL2(dimension)
        else:
            # For larger datasets, use IVF index
            nlist = min(int(np.sqrt(len(embeddings))), 1000)
            quantizer = faiss.IndexFlatL2(dimension)
            index = faiss.IndexIVFFlat(quantizer, dimension, nlist)
            # Train the index
            index.train(embeddings_array)
        
        index.add(embeddings_array)
        faiss.write_index(index, "app/data/faiss.index")
        logger.info(f"FAISS index saved (type: {type(index).__name__}, dimension: {dimension})")
    
    # Step 6: Save metadata files
    with open("app/data/id_map.json", "w") as f:
        json.dump(id_map, f)
    
    with open("app/data/movie_metadata.json", "w") as f:
        json.dump(movie_metadata, f)
    
    logger.info("✅ Index built successfully!")
    logger.info(f"   - {len(embeddings)} movies indexed")
    logger.info(f"   - Embedding model: {model}")
    logger.info(f"   - Files saved in app/data/")
    logger.info(f"     * movies.npy: embeddings matrix")
    logger.info(f"     * id_map.json: TMDB ID to matrix position mapping")
    logger.info(f"     * movie_metadata.json: movie metadata")
    if use_faiss:
        logger.info(f"     * faiss.index: FAISS search index")
    
    # Cleanup checkpoints
    cleanup_checkpoints()

# Remove the old functions that are no longer needed
# create_movie_embedding and load_movie_data are replaced by the new implementation

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Build movie embeddings index from TMDB data")
    parser.add_argument("--max-pages", type=int, default=10, 
                       help="Maximum pages to fetch from TMDB popular movies (default: 10)")
    parser.add_argument("--model", type=str, default="text-embedding-3-small",
                       help="OpenAI embedding model to use (default: text-embedding-3-small)")
    parser.add_argument("--no-faiss", action="store_true",
                       help="Skip building FAISS index")
    
    args = parser.parse_args()
    
    build_index(
        max_pages=args.max_pages,
        model=args.model,
        use_faiss=not args.no_faiss
    )