"""Common module for placements lambda functions."""

import csv

import smart_open
from sosmodels import streams_vendor_placements
from sosmodels import streams_track_placements

from sosmodels.active_tables import get_incremented_connection

import consts  # noqa
import util


def load(s3_key_url, placements_type):
    """Load placements data to passed model.

    Args:
        s3_key_url (str): S3 url for loading file.
        placements_type (str): Placements type (vendor or track).
    """
    store = util.identify_retailer(s3_key_url)

    PLACEMENTS_MODEL_MAPPING = {
        consts.VENDOR: {
            consts.SPOTIFY:
                streams_vendor_placements.StreamsVendorPlacementsSpotify,
            consts.APPLE_MUSIC:
                streams_vendor_placements.StreamsVendorPlacementsAppleMusic
        },
        consts.TRACK: {
            consts.SPOTIFY:
                streams_track_placements.StreamsTrackPlacementsSpotify,
            consts.APPLE_MUSIC:
                streams_track_placements.StreamsTrackPlacementsAppleMusic
        }
    }
    try:
        placements_model = PLACEMENTS_MODEL_MAPPING[placements_type][store]
    except KeyError:
        raise KeyError('Incorrect placements type.')
    # TODO: refactor within DATA-3267
    if placements_type == consts.VENDOR:
        columns = [
            'vendor_placement_key', 'order_number', 'streams', 'storeid',
            'playlist_title', 'playlist_author', 'playlist_link',
            'track_title', 'upc', 'isrc', 'artist', 'date_added', 'start_date',
            'end_date']
    elif placements_type == consts.TRACK:
        columns = [
            'track_placement_key', 'order_number', 'streams',
            'storeid', 'playlist_title', 'playlist_author',
            'playlist_link', 'date_added', 'start_date', 'end_date']
    else:
        raise TypeError('Incorrect placements type.')

    placements_model._get_connection = get_incremented_connection

    with placements_model.batch_write() as batch:
        with smart_open.smart_open(s3_key_url) as lines:
            for placements in csv.reader(lines, dialect='escaped'):
                assert len(columns) == len(placements)
                row = dict(zip(columns, placements))
                item_data = {
                    'order_number': int(row['order_number']),
                    'streams': int(row['streams']),
                    'storeid': int(row['storeid']),
                    'playlist_title': row['playlist_title'].decode('utf8'),
                    'playlist_author': row['playlist_author'].decode('utf8'),
                    'playlist_link': row['playlist_link'].decode('utf8'),
                    'date_added': row['date_added'].decode('utf8'),
                    'start_date': row['start_date'].decode('utf8'),
                    'end_date': row['end_date'].decode('utf8'),
                }
                if placements_type == consts.VENDOR:
                    item_data['vendor_placement_key'] = row[
                        'vendor_placement_key']
                    item_data['upc'] = row['upc'].decode('utf8')
                    item_data['track_title'] = row[
                        'track_title'].decode('utf8')
                    item_data['isrc'] = row['isrc'].decode('utf8')
                    item_data['artist'] = row['artist'].decode('utf8')
                elif placements_type == consts.TRACK:
                    item_data['track_placement_key'] = row[
                        'track_placement_key']
                item = placements_model(**item_data)
                batch.save(item)