mtgonline/backend/scripts/load_mtgjson_data.py

#!/usr/bin/env python3
"""
MTGJSON Data Loader

Downloads and loads MTGJSON data into the PostgreSQL database.
Handles AllPrintings.psql, JSON files, and image data extraction.

Usage:
    python load_mtgjson_data.py
"""

import asyncio
import json
import gzip
import logging
import os
import sys
import tempfile
from pathlib import Path
from urllib.request import urlretrieve

# Add parent directory to path for imports
sys.path.append(str(Path(__file__).parent.parent))

from app.core.settings import get_settings
from app.core.database import mtg_engine, mtg_async_session
from app.models.mtg_models import MtgSet, MtgCard

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# MTGJSON download URLs
MTGJSON_BASE_URL = "https://mtgjson.com/api/5x"
MTGJSON_FILES = {
    "AllPrintings.psql.gz": "AllPrintings.psql.gz",
    "AllSetFiles.zip": "AllSetFiles.zip",
    "AllDeckFiles.zip": "AllDeckFiles.zip",
    "AllIdentifiers.json.gz": "AllIdentifiers.json.gz",
    "CardTypes.json.gz": "CardTypes.json.gz",
    "DeckList.json.gz": "DeckList.json.gz",
    "Keywords.json.gz": "Keywords.json.gz",
    "SetList.json.gz": "SetList.json.gz",
}


async def load_sets_from_json():
    """Load set metadata from AllSetFiles.zip or SetList.json.gz"""
    logger.info("Loading set metadata...")
    
    # Try AllSetFiles first (has more complete data)
    set_files_dir = Path("/app/data/mtgjson/allsetfiles")
    if not set_files_dir.exists():
        # Try SetList.json.gz as fallback
        setlist_path = Path("/app/data/mtgjson/SetList.json.gz")
        if setlist_path.exists():
            logger.info("Loading from SetList.json.gz")
            with gzip.open(setlist_path, 'rt', encoding='utf-8') as f:
                set_list = json.load(f)
            
            settings = get_settings()
            async with mtg_async_session() as session:
                for set_data in set_list:
                    # Get image URL from setCode mapping if available
                    image_url = None
                    if 'image' in set_data:
                        image_url = set_data['image'].get('png', set_data['image'].get('svg'))
                    
                    existing = await session.execute(
                        MtgSet.__table__.select().where(MtgSet.code == set_data['code'])
                    )
                    
                    if existing.first():
                        # Update existing
                        await session.execute(
                            MtgSet.__table__.update()
                            .where(MtgSet.code == set_data['code'])
                            .values(
                                name=set_data.get('name'),
                                type=set_data.get('type'),
                                release_date=set_data.get('releaseDate'),
                                base_set_size=set_data.get('baseSetSize'),
                                total_size=set_data.get('totalSize'),
                                icon_svg_url=set_data.get('iconSvgUri'),
                                image=image_url,
                                updated_at=asyncio.coroutines.utcnow() if hasattr(asyncio.coroutines, 'utcnow') else None
                            )
                        )
                    else:
                        # Insert new
                        set_obj = MtgSet(
                            code=set_data['code'],
                            name=set_data.get('name'),
                            type=set_data.get('type'),
                            release_date=set_data.get('releaseDate'),
                            base_set_size=set_data.get('baseSetSize'),
                            total_size=set_data.get('totalSize'),
                            icon_svg_url=set_data.get('iconSvgUri'),
                            image=image_url,
                            created_at=asyncio.coroutines.utcnow() if hasattr(asyncio.coroutines, 'utcnow') else None
                        )
                        session.add(set_obj)
                
                await session.commit()
                logger.info(f"Loaded set metadata from {len(set_list)} sets")
        else:
            logger.warning("No set metadata files found")
            return
    
    # Process AllSetFiles directory
    count = 0
    settings = get_settings()
    async with mtg_async_session() as session:
        for json_file in sorted(set_files_dir.glob("*.json")):
            with open(json_file, 'r', encoding='utf-8') as f:
                set_data = json.load(f)
            
            if 'data' in set_data:
                set_data = set_data['data']
            
            image_url = set_data.get('image', {}).get('png', set_data.get('image', {}).get('svg'))
            
            existing = await session.execute(
                MtgSet.__table__.select().where(MtgSet.code == set_data['code'])
            )
            
            if existing.first():
                await session.execute(
                    MtgSet.__table__.update()
                    .where(MtgSet.code == set_data['code'])
                    .values(
                        name=set_data.get('name'),
                        type=set_data.get('type'),
                        release_date=set_data.get('releaseDate'),
                        base_set_size=set_data.get('baseSetSize'),
                        total_size=set_data.get('totalSize'),
                        icon_svg_url=set_data.get('iconSvgUri'),
                        image=image_url
                    )
                )
            else:
                set_obj = MtgSet(
                    code=set_data['code'],
                    name=set_data.get('name'),
                    type=set_data.get('type'),
                    release_date=set_data.get('releaseDate'),
                    base_set_size=set_data.get('baseSetSize'),
                    total_size=set_data.get('totalSize'),
                    icon_svg_url=set_data.get('iconSvgUri'),
                    image=image_url
                )
                session.add(set_obj)
            
            count += 1
        
        await session.commit()
        logger.info(f"Loaded {count} sets from AllSetFiles")


async def load_cards_from_psql():
    """Load cards from AllPrintings.psql"""
    logger.info("Loading cards from AllPrintings.psql...")
    
    psql_path = Path("/app/data/mtgjson/AllPrintings.psql")
    if not psql_path.exists():
        logger.warning("AllPrintings.psql not found")
        return
    
    # Parse PSQL file to extract INSERT statements
    # This is a simplified parser - in production you'd use a proper PSQL parser
    cards_data = []
    
    with open(psql_path, 'r', encoding='utf-8') as f:
        current_card = {}
        in_insert = False
        
        for line in f:
            line = line.strip()
            
            if line.startswith('COPY public.mtgjson_card'):
                # Header line - skip
                continue
            
            if line == '\\.':
                # End of COPY command
                in_insert = False
                continue
            
            if in_insert:
                # Parse CSV line
                fields = line.split('\t')
                if len(fields) > 10:
                    try:
                        card = {
                            'id': fields[0],
                            'name': fields[1],
                            'manaCost': fields[2],
                            'type': fields[3],
                            'text': fields[4],
                            'power': fields[5],
                            'toughness': fields[6],
                            'rarity': fields[7],
                            'layout': fields[8],
                            'artist': fields[9],
                            'flavor': fields[10] if len(fields) > 10 else '',
                            'set': fields[11] if len(fields) > 11 else '',
                            'number': fields[12] if len(fields) > 12 else '',
                            'identifiers': fields[13] if len(fields) > 13 else '{}',
                            'images': fields[14] if len(fields) > 14 else '{}',
                            'updatedAt': fields[15] if len(fields) > 15 else '',
                        }
                        cards_data.append(card)
                    except (ValueError, IndexError):
                        continue
            
            if line.startswith('INSERT INTO public.mtgjson_card'):
                in_insert = True
    
    logger.info(f"Parsed {len(cards_data)} cards from PSQL file")
    
    # Update cards in database with parsed data
    if cards_data:
        settings = get_settings()
        async with mtg_async_session() as session:
            # First, get all set codes to create sets
            set_codes = set(c['set'] for c in cards_data if c['set'])
            
            for set_code in set_codes:
                existing = await session.execute(
                    MtgSet.__table__.select().where(MtgSet.code == set_code)
                )
                if not existing.first():
                    # Create placeholder set
                    set_obj = MtgSet(
                        code=set_code,
                        name=f"Set {set_code}",
                        created_at=asyncio.coroutines.utcnow() if hasattr(asyncio.coroutines, 'utcnow') else None
                    )
                    session.add(set_obj)
            
            await session.flush()
            
            # Now load cards
            for card_data in cards_data:
                # Get set_id
                set_result = await session.execute(
                    MtgSet.__table__.select().where(MtgSet.code == card_data['set'])
                )
                set_obj = set_result.first()
                
                if not set_obj:
                    continue
                
                existing = await session.execute(
                    MtgCard.__table__.select()
                    .where(MtgCard.name == card_data['name'])
                    .where(MtgCard.set_id == set_obj.id)
                )
                
                if existing.first():
                    # Update existing
                    await session.execute(
                        MtgCard.__table__.update()
                        .where(MtgCard.name == card_data['name'])
                        .where(MtgCard.set_id == set_obj.id)
                        .values(
                            mana_cost=card_data['manaCost'],
                            type_line=card_data['type'],
                            oracle_text=card_data['text'],
                            power=card_data['power'],
                            toughness=card_data['toughness'],
                            rarity=card_data['rarity'],
                            layout=card_data['layout'],
                            artist=card_data['artist'],
                            flavor_text=card_data['flavor'],
                            numbers=card_data['number'],
                            identifiers=card_data['identifiers'],
                            images=card_data['images'],
                            image=card_data.get('image'),
                            updated_at=asyncio.coroutines.utcnow() if hasattr(asyncio.coroutines, 'utcnow') else None
                        )
                    )
                else:
                    # Insert new card
                    card_obj = MtgCard(
                        name=card_data['name'],
                        set_id=set_obj.id,
                        mana_cost=card_data['manaCost'],
                        type_line=card_data['type'],
                        oracle_text=card_data['text'],
                        power=card_data['power'],
                        toughness=card_data['toughness'],
                        rarity=card_data['rarity'],
                        layout=card_data['layout'],
                        artist=card_data['artist'],
                        flavor_text=card_data['flavor'],
                        numbers=card_data['number'],
                        identifiers=card_data['identifiers'],
                        images=card_data['images'],
                        image=card_data.get('image')
                    )
                    session.add(card_obj)
            
            await session.commit()
            logger.info("Cards loaded successfully")


async def load_identifiers():
    """Load card identifiers from AllIdentifiers.json.gz"""
    logger.info("Loading identifiers...")
    
    identifiers_path = Path("/app/data/mtgjson/AllIdentifiers.json.gz")
    if not identifiers_path.exists():
        logger.warning("AllIdentifiers.json.gz not found")
        return
    
    with gzip.open(identifiers_path, 'rt', encoding='utf-8') as f:
        identifiers = json.load(f)
    
    logger.info(f"Loaded {len(identifiers)} identifiers")


async def load_deck_list():
    """Load deck list metadata from DeckList.json.gz"""
    logger.info("Loading deck list...")
    
    deck_list_path = Path("/app/data/mtgjson/DeckList.json.gz")
    if not deck_list_path.exists():
        logger.warning("DeckList.json.gz not found")
        return
    
    with gzip.open(deck_list_path, 'rt', encoding='utf-8') as f:
        deck_list = json.load(f)
    
    logger.info(f"Loaded {len(deck_list)} deck list entries")


async def load_keywords():
    """Load card keywords from Keywords.json.gz"""
    logger.info("Loading keywords...")
    
    keywords_path = Path("/app/data/mtgjson/Keywords.json.gz")
    if not keywords_path.exists():
        logger.warning("Keywords.json.gz not found")
        return
    
    with gzip.open(keywords_path, 'rt', encoding='utf-8') as f:
        keywords = json.load(f)
    
    logger.info(f"Loaded {len(keywords)} keywords")


async def main():
    """Main entry point"""
    logger.info("Starting MTGJSON data loader...")
    
    # Ensure data directory exists
    data_dir = Path("/app/data/mtgjson")
    data_dir.mkdir(parents=True, exist_ok=True)
    
    # Load data in order
    await load_sets_from_json()
    await load_cards_from_psql()
    await load_identifiers()
    await load_deck_list()
    await load_keywords()
    
    logger.info("MTGJSON data loading complete!")
    
    # Print summary
    async with mtg_async_session() as session:
        from sqlalchemy import text
        
        result = await session.execute(text("SELECT COUNT(*) FROM mtg_sets"))
        set_count = result.scalar()
        
        result = await session.execute(text("SELECT COUNT(*) FROM mtg_cards"))
        card_count = result.scalar()
        
        logger.info(f"Database summary:")
        logger.info(f"  Sets: {set_count}")
        logger.info(f"  Cards: {card_count}")


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