q3stats/q3stats/lib/scripts/import_games.py

# -*- coding: utf-8 -*-
# Copyright (c) 2017 Tomek Wójcik <tomek@bthlabs.pl>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#

"""
q3stats.lib.scripts.import_games
================================

This module contains the Q3A game importer.
"""

from argparse import ArgumentParser
import datetime
import codecs
import logging
import os

import lxml.etree

from q3stats.lib.config import get_config
from q3stats.lib.scripts import utils
from q3stats.models import Game, Score

__all__ = ['script_main']

XML_STAT_TO_MODEL_FIELD = {
    'Score': 'score',
    'Kills': 'kills',
    'Deaths': 'deaths',
    'Suicides': 'suicides',
    'Net': 'net',
    'DamageGiven': 'damage_given',
    'DamageTaken': 'damage_taken',
    'HealthTotal': 'total_health',
    'ArmorTotal': 'total_armor',
}

XML_STAT_POWERUP_NAMES = ['Quad']


def _find_stats_files(stats_path):
    """Finds Q3A stats files in *stats_path*."""
    stats_files = []
    for root, dirs, files in os.walk(stats_path):
        for filename in files:
            _, ext = os.path.splitext(filename)
            if ext.lower() == '.xml':
                stats_files.append(os.path.join(root, filename))

    return stats_files


def _read_match(file_path):
    """Reads a Q3A match file from *file_path*."""
    with codecs.open(file_path, 'r', 'utf-8') as stats_f:
        stats_xml = lxml.etree.parse(stats_f)

    xmlroot = stats_xml.getroot()
    if xmlroot.tag != 'match':
        return None

    return xmlroot


def _read_game(match):
    """Creates and returns an instance of :py:class:`q3stats.models.game.Game`
    filled with game data from *match*."""
    ts = datetime.datetime.strptime(
        match.get('datetime'), '%Y/%m/%d %H:%M:%S'
    )

    is_team = False
    if match.get('isTeamGame') == 'true':
        is_team = True

    game = Game(
        map=match.get('map').upper(),
        date=ts.date(),
        time=ts.time(),
        attrs={
            'type': match.get('type'),
            'team_game': is_team,
            'duration': int(match.get('duration')),
        }
    )

    game.update_uuid()

    return game


def _read_scores(match):
    """Reads game data from *match* and extracts player scores.

    Returns an array of :py:class:`q3stats.models.score.Score` objects and
    the actual fraglimit (aka the highest score)."""
    scores = []

    current_score = None
    fraglimit = 0
    for element in match.iter():
        if element.tag == 'player':
            if current_score is not None:
                if current_score.score > fraglimit:
                    fraglimit = current_score.score

                scores.append(current_score)

            current_score = Score(player=element.get('name'))
            current_score.weapons = {}
            current_score.items = {}
            current_score.powerups = {}
        elif element.tag == 'stat':
            name = element.get('name')
            if name in XML_STAT_TO_MODEL_FIELD:
                setattr(
                    current_score, XML_STAT_TO_MODEL_FIELD[name],
                    int(element.get('value'))
                )
        elif element.tag == 'weapon':
            current_score.weapons[element.get('name')] = {
                'hits': int(element.get('hits')),
                'shots': int(element.get('shots')),
                'kills': int(element.get('kills')),
            }
        elif element.tag == 'item':
            name = element.get('name')
            if name in XML_STAT_POWERUP_NAMES:
                current_score.powerups[name] = [
                    int(element.get('pickups')), int(element.get('time'))
                ]
            else:
                current_score.items[name] = int(element.get('pickups'))

    if current_score is not None:
        if current_score.score > fraglimit:
            fraglimit = current_score.score

        scores.append(current_score)

    return scores, fraglimit


def _read_stats_file(file_path):
    """Reads stats file from *file_path*.

    Returns an instance of :py:class:`q3stats.models.game.Game` with
    associated instances of :py:class:`q3stats.models.score.Score`."""
    match = _read_match(file_path)
    if match is None:
        raise RuntimeError('Invalid stats file: %s' % file_path)

    game = _read_game(match)

    scores, fraglimit = _read_scores(match)
    game.scores.extend(scores)
    game.fraglimit = fraglimit

    return game


def _get_argument_parser():
    """Creates and returns an instance of ``ArgumentParser`` configured for
    the script."""
    parser = ArgumentParser(description='Import Q3A stats')
    parser.add_argument('stats_path', help='path to stats directory')
    parser.add_argument(
        '-v', '--verbose', action='store_true', default=False,
        help='be verbose'
    )

    return parser


def script_main(argv, config=None):
    """Script entry point."""
    result = utils.RET_OK

    parser = _get_argument_parser()
    args = parser.parse_args(args=argv)

    logger = utils.get_logger('import_games', args)

    try:
        if config is None:
            config = get_config()

        stats_files = _find_stats_files(args.stats_path)
        logger.debug('Found %d stat files.' % len(stats_files))

        with utils.db_session(config) as session:
            for stats_file in stats_files:
                logger.debug('Processing file: %s' % stats_file)
                game = _read_stats_file(stats_file)

                current_game = session.query(Game).\
                    filter_by(uuid=game.uuid).\
                    first()

                if current_game is not None:
                    logger.debug('Skipping!')
                else:
                    session.add(game)

            session.commit()
    except Exception as exc:
        logger.error('Unhandled exception', exc_info=True)
        result = utils.RET_ERROR
    finally:
        return result