"""
Class StreamToS3
================

This class will write a CSV to a stream, and write the stream
to an S3 bucket. It will empty the stream as it writes to S3.
If the file already exists in S3, it will replace it.
"""


import csv
import io

import smart_open


class StreamToS3():

    def __init__(self, key, bucket, row_dict_keys, column_headings=None):
        """Set class variables.

        If your actual display column headers for the CSV file are
        different than your row data key values, and you want a header on
        your columns, pass the optional 'column_headings' param.

        Args:
            key (str): S3 key name
            bucket (str): S3 bucket name
            row_dict_keys (list): list of keys in your csv row dictionaries
                in the order you wish your columns to appear.
            column_headings (list): optional list of display headings
                for columns in CSV file. (Make sure you pass in same order as
                your row_dict_keys for proper mapping.)
        """

        self.key = key
        self.bucket = bucket
        self.dict_keys = row_dict_keys
        self.column_headings = column_headings

    def __enter__(self):
        """Creates a context manager for StreamToS3.

        The context manager initializes the stream and opens file on S3.

        Return:
            StreamToS3: the instantiated object StreamToS3.
        """

        path = 's3://{}/{}'.format(self.bucket, self.key)
        self.s3writer = smart_open.smart_open(path, 'wb')
        self.stream = io.StringIO()
        return self

    def __exit__(self, type, value, traceback):
        """Closes S3 file and memory stream."""

        self.s3writer.close()
        self.stream.close()

    def set_csv_writer(self, column_headings, clear_stream=True):
        """Give the writer object a stream and set for class use.

        Args:
            column_headings (list): required by the CSV writer. The
            dict keys in the order you wish them to map to columns.
            clear_stream (bool): Get a new stream (ie. clear memory)
        """

        if clear_stream:
            self.stream = io.StringIO()

        self.csv_writer = csv.DictWriter(
            self.stream, fieldnames=column_headings, extrasaction='ignore')

    @property
    def header(self):
        """Get the headers of the csv

        Return:
            list: the header list.
        """

        return self.column_headings or self.dict_keys

    def write_column_headers(self):
        """Write a header row with column names to the CSV file on S3."""

        self.set_csv_writer(self.header, clear_stream=False)
        self.csv_writer.writeheader()
        self.s3writer.write(self.stream.getvalue())

    def write_batch_rows(self, rows):
        """Write a batch of CSV rows to the S3 file.

        Writes a batch of CSV rows to Stream and then to S3 file. Stream
        memory clears at end of each batch.

        Args:
            rows (list): list of dictionaries where each dict is a new CSV row.
        """

        self.set_csv_writer(self.dict_keys)
        for row in rows:
            if ('upc' in row):
                row['upc'] = '="' + row['upc'] + '"'
            if ('iodaUpc' in row):
                row['iodaUpc'] = '="' + row['iodaUpc'] + '"'
            self.csv_writer.writerow(row)

        self.s3writer.write(self.stream.getvalue())