mtgonline/backend/scripts/download_mtgjson.py

"""
MTGJSON Download Script

Downloads MTGJSON files from the official API and processes them into PostgreSQL.
Designed to run on container init and when refresh is triggered via API.

## Features
- Parallel downloads with retry logic
- Progress tracking and logging
- Automatic decompression
- Verification and integrity checks
- Database upsert operations

## Usage
```bash
# Download and process all files
python download_mtgjson.py --all

# Download specific files
python download_mtgjson.py --files AllPrintings.psql AllIdentifiers.json

# Force re-download
python download_mtgjson.py --all --force

# Dry run (show what would be downloaded)
python download_mtgjson.py --all --dry-run
```
"""

import asyncio
import aiohttp
import gzip
import zipfile
import shutil
import tempfile
import logging
from pathlib import Path
from datetime import datetime
from typing import Optional, Dict, List, Tuple
from dataclasses import dataclass
from enum import Enum

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('mtgjson_download.log', mode='a')
    ]
)
logger = logging.getLogger(__name__)


class DownloadStatus(Enum):
    """Status of download operations."""
    PENDING = "pending"
    DOWNLOADING = "downloading"
    DECOMPRESSING = "decompressing"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"


@dataclass
class MTGFile:
    """Represents an MTGJSON file to download."""
    name: str
    filename: str
    compressed: bool = True
    description: str = ""
    
    def __post_init__(self):
        if not self.description:
            self.description = f"{self.name} file"


class MTGJSONDownloader:
    """Download and process MTGJSON files from the official API."""
    
    # Base URL for MTGJSON API
    BASE_URL = "https://mtgjson.com/api/v5/"
    
    # Configuration
    MAX_RETRIES = 3
    RETRY_DELAY = 5  # seconds
    CHUNK_SIZE = 8192  # bytes per chunk for progress tracking
    TIMEOUT = aiohttp.ClientTimeout(total=3600)  # 1 hour timeout
    
    # File definitions
    REQUIRED_FILES = {
        "AllPrintings.psql": MTGFile(
            name="AllPrintings",
            filename="AllPrintings.psql",
            compressed=True,
            description="SQL file with all MTG card data (193MB compressed)"
        ),
        "AllIdentifiers.json": MTGFile(
            name="AllIdentifiers",
            filename="AllIdentifiers.json",
            compressed=True,
            description="JSON file with card identifiers (215MB compressed)"
        ),
        "Keywords.json": MTGFile(
            name="Keywords",
            filename="Keywords.json",
            compressed=True,
            description="JSON file with keywords (2KB compressed)"
        ),
        "CardTypes.json": MTGFile(
            name="CardTypes",
            filename="CardTypes.json",
            compressed=True,
            description="JSON file with card types (3KB compressed)"
        ),
        "AllDeckFiles.zip": MTGFile(
            name="AllDeckFiles",
            filename="AllDeckFiles.zip",
            compressed=False,
            description="ZIP archive with deck files (246MB)"
        ),
    }
    
    def __init__(self, data_dir: Path, force: bool = False, dry_run: bool = False):
        """
        Initialize the downloader.
        
        Args:
            data_dir: Directory to store downloaded files
            force: Force re-download even if files exist
            dry_run: Show what would be done without actually doing it
        """
        self.data_dir = Path(data_dir)
        self.force = force
        self.dry_run = dry_run
        
        # Create data directory if it doesn't exist
        self.data_dir.mkdir(parents=True, exist_ok=True)
        
        # Track download status
        self.download_status: Dict[str, DownloadStatus] = {}
        for filename in self.REQUIRED_FILES:
            self.download_status[filename] = DownloadStatus.PENDING
        
        logger.info(f"MTGJSON Downloader initialized")
        logger.info(f"Data directory: {self.data_dir}")
        logger.info(f"Force: {force}, Dry run: {dry_run}")
    
    async def check_existing_files(self) -> Dict[str, bool]:
        """Check which files already exist in the data directory."""
        existing = {}
        for filename in self.REQUIRED_FILES:
            file_path = self.data_dir / filename
            if file_path.exists():
                size_mb = file_path.stat().st_size / (1024 * 1024)
                logger.info(f"Found existing file: {filename} ({size_mb:.1f} MB)")
                existing[filename] = True
            else:
                logger.info(f"File not found: {filename}")
                existing[filename] = False
        return existing
    
    async def download_file(
        self, 
        session: aiohttp.ClientSession,
        filename: str,
        progress_callback: Optional[callable] = None
    ) -> Tuple[bool, str]:
        """
        Download a single file with retry logic.
        
        Args:
            session: aiohttp session
            filename: Name of the file to download
            progress_callback: Optional callback for progress updates
            
        Returns:
            Tuple of (success, message)
        """
        file_info = self.REQUIRED_FILES[filename]
        file_path = self.data_dir / filename
        
        for attempt in range(1, self.MAX_RETRIES + 1):
            try:
                logger.info(f"Downloading {filename} (attempt {attempt}/{self.MAX_RETRIES})")
                self.download_status[filename] = DownloadStatus.DOWNLOADING
                
                # Set up URL
                url = f"{self.BASE_URL}{filename}"
                
                # Download with progress
                async with session.get(url, timeout=self.TIMEOUT) as response:
                    if response.status != 200:
                        error_msg = f"HTTP {response.status} for {filename}"
                        logger.error(error_msg)
                        
                        if attempt < self.MAX_RETRIES:
                            await asyncio.sleep(self.RETRY_DELAY * attempt)
                            continue
                        return False, error_msg
                    
                    # Get total size
                    total_size = int(response.headers.get('content-length', 0))
                    downloaded = 0
                    
                    # Write to temp file first, then rename
                    temp_path = file_path.with_suffix(file_path.suffix + '.tmp')
                    
                    with open(temp_path, 'wb') as f:
                        async for chunk in response.content.iter_chunked(self.CHUNK_SIZE):
                            f.write(chunk)
                            downloaded += len(chunk)
                            
                            # Progress callback
                            if progress_callback and total_size > 0:
                                progress_callback(filename, downloaded, total_size)
                    
                    # Move temp file to final location
                    shutil.move(str(temp_path), str(file_path))
                    
                    # Decompress if needed
                    if file_info.compressed:
                        await self._decompress_file(file_path)
                    
                    logger.info(f"Successfully downloaded {filename}")
                    self.download_status[filename] = DownloadStatus.COMPLETED
                    return True, "Download completed"
                    
            except Exception as e:
                error_msg = f"Error downloading {filename}: {str(e)}"
                logger.error(error_msg)
                
                # Clean up temp file if it exists
                temp_path = file_path.with_suffix(file_path.suffix + '.tmp')
                if temp_path.exists():
                    temp_path.unlink()
                
                if attempt < self.MAX_RETRIES:
                    await asyncio.sleep(self.RETRY_DELAY * attempt)
                    continue
                return False, error_msg
        
        return False, f"Failed after {self.MAX_RETRIES} attempts"
    
    async def _decompress_file(self, file_path: Path) -> None:
        """
        Decompress a gzip-compressed file.
        
        Args:
            file_path: Path to the compressed file
        """
        self.download_status[file_path.name] = DownloadStatus.DECOMPRESSING
        
        if file_path.suffix == '.gz':
            decompressed_path = file_path.with_suffix('')
            
            logger.info(f"Decompressing {file_path.name} -> {decompressed_path.name}")
            
            with gzip.open(file_path, 'rb') as f_in:
                with open(decompressed_path, 'wb') as f_out:
                    shutil.copyfileobj(f_in, f_out)
            
            # Remove compressed file
            file_path.unlink()
            
            logger.info(f"Decompression complete: {decompressed_path.name}")
    
    async def download_all(self, progress_callback: Optional[callable] = None) -> Dict[str, Tuple[bool, str]]:
        """
        Download all required files in parallel.
        
        Args:
            progress_callback: Optional callback for progress updates
            
        Returns:
            Dictionary of filename -> (success, message)
        """
        results = {}
        
        async with aiohttp.ClientSession() as session:
            # Check existing files
            existing = await self.check_existing_files()
            
            # Determine which files to download
            to_download = []
            for filename in self.REQUIRED_FILES:
                if self.force or not existing[filename]:
                    to_download.append(filename)
                else:
                    logger.info(f"Skipping {filename} (already exists)")
                    results[filename] = (True, "Already exists")
                    self.download_status[filename] = DownloadStatus.COMPLETED
            
            if not to_download:
                logger.info("All files already exist, nothing to download")
                return results
            
            # Download files in parallel
            logger.info(f"Downloading {len(to_download)} files in parallel...")
            
            tasks = []
            for filename in to_download:
                task = self.download_file(session, filename, progress_callback)
                tasks.append(task)
            
            download_results = await asyncio.gather(*tasks)
            
            # Collect results
            for filename, result in zip(to_download, download_results):
                results[filename] = result
        
        return results
    
    def print_summary(self, results: Dict[str, Tuple[bool, str]]) -> None:
        """Print a summary of download results."""
        logger.info("\n" + "="*60)
        logger.info("DOWNLOAD SUMMARY")
        logger.info("="*60)
        
        success_count = 0
        for filename, (success, message) in results.items():
            status_icon = "✓" if success else "✗"
            logger.info(f"  {status_icon} {filename}: {message}")
            if success:
                success_count += 1
        
        logger.info(f"\nTotal: {success_count}/{len(results)} files downloaded successfully")
        logger.info("="*60 + "\n")
    
    async def verify_files(self) -> Tuple[bool, List[str]]:
        """
        Verify that all required files exist and are valid.
        
        Returns:
            Tuple of (all_valid, list_of_errors)
        """
        errors = []
        
        for filename, file_info in self.REQUIRED_FILES.items():
            file_path = self.data_dir / filename
            
            if not file_path.exists():
                errors.append(f"Missing required file: {filename}")
                continue
            
            # Check file size (basic sanity check)
            size = file_path.stat().st_size
            if size == 0:
                errors.append(f"Empty file: {filename}")
                continue
            
            # Verify JSON files
            if filename.endswith('.json'):
                try:
                    import json
                    with open(file_path, 'r') as f:
                        json.load(f)
                except Exception as e:
                    errors.append(f"Invalid JSON in {filename}: {str(e)}")
            
            # Verify ZIP files
            if filename.endswith('.zip'):
                try:
                    with zipfile.ZipFile(file_path, 'r') as zf:
                        zf.testzip()
                except Exception as e:
                    errors.append(f"Invalid ZIP file {filename}: {str(e)}")
            
            # Verify SQL files
            if filename.endswith('.psql'):
                # Basic check - file should not be empty and should have some SQL content
                with open(file_path, 'r') as f:
                    content = f.read(1024)
                    if not any(kw in content.upper() for kw in ['INSERT', 'CREATE', 'BEGIN']):
                        errors.append(f"File {filename} doesn't appear to contain SQL")
        
        all_valid = len(errors) == 0
        return all_valid, errors


async def main():
    """Main entry point."""
    import argparse
    
    parser = argparse.ArgumentParser(description="Download MTGJSON files")
    parser.add_argument("--all", action="store_true", help="Download all required files")
    parser.add_argument("--files", nargs="+", help="Specific files to download")
    parser.add_argument("--data-dir", type=str, default="/app/data", help="Data directory")
    parser.add_argument("--force", action="store_true", help="Force re-download")
    parser.add_argument("--dry-run", action="store_true", help="Show what would be done")
    parser.add_argument("--verify", action="store_true", help="Verify files after download")
    
    args = parser.parse_args()
    
    data_dir = Path(args.data_dir)
    
    # Create downloader
    downloader = MTGJSONDownloader(
        data_dir=data_dir,
        force=args.force,
        dry_run=args.dry_run
    )
    
    # Determine which files to download
    if args.files:
        files_to_download = args.files
    elif args.all:
        files_to_download = list(MTGJSONDownloader.REQUIRED_FILES.keys())
    else:
        logger.info("No files specified. Use --all or --files <filename>")
        return 1
    
    # Download files
    results = await downloader.download_all()
    
    # Print summary
    downloader.print_summary(results)
    
    # Verify if requested
    if args.verify:
        all_valid, errors = await downloader.verify_files()
        if all_valid:
            logger.info("All files verified successfully!")
        else:
            logger.error("Verification failed:")
            for error in errors:
                logger.error(f"  - {error}")
            return 1
    
    return 0


if __name__ == "__main__":
    exit_code = asyncio.run(main())
    exit(exit_code)