"""Producer script for filling the Kinesis stream with reports data."""

import csv
from datetime import datetime
import glob
import gzip
import os
import sys
import tempfile

from pymysql import cursors
from pymysql import err
from snowflake.connector.errors import Error as SnowflakeError

from accounting import config
from accounting import const
from accounting import util
from accounting.config import log
from accounting.connectors import mysql
from accounting.connectors import snowflake
from accounting.models import sql
from accounting.models import sql_fp
from accounting.models import statement_db

csv.field_size_limit(10_000_000)


def export_csv(file_name):
    """Read the data from stmt-db and save as CSV.

    Args:
        file_name (str): csv file name to save

    Returns:
        int: number of rows exported

    """
    conn = mysql.get_connection(cursorclass=cursors.SSDictCursor)
    rows_processed = 0
    use_mrr = config.USE_MRR
    # TODO: Remove this once fingerprinting is live
    fieldnames = [const.ISRC, const.TERRITORY] if use_mrr \
        else [const.SERVICE, const.DATE_COLUMN_NAME,
              const.ISRC, const.TERRITORY]
    try:
        with open(file_name, 'w') as f:
            writer = csv.DictWriter(
                f, fieldnames=fieldnames)
            with conn.cursor() as cursor:
                log.info('Batch received.')
                statement_db.get_batch(cursor, config.MAX_RECORDS, 0)
                for row in cursor:
                    writer.writerow(row)
                    rows_processed += 1
                    if rows_processed >= config.MAX_RECORDS:
                        break
                    if rows_processed % 10000 == 0:
                        log.info(f'{rows_processed} written.')
    except err.MySQLError:
        log.exception('Failed to read statement data.')
        sys.exit(1)
    finally:
        conn.close()
    return rows_processed


def process_snowflake(file_name, output_dir, conn=None):
    """Import CSV data into snowflake and get the corresponding TUIDs.

    The steps are:
    1. Import the CSV into Snowflake tmp table
    2. Update tmp table with tuid and internal conflict information
    from facts.prod.registry
    3. Export an updated tmp table to CSV
    4. Download the exported CSV file

    Args:
        file_name (str): csv file to process
        output_dir (str): output directory to save the processed data

    """
    if conn is None:
        conn = snowflake.get_snowflake_connection()
    try:
        conn.cursor().execute(sql.CREATE_TMP_TABLE)
        conn.cursor().execute(sql.PUT_CSV.format(file_name=file_name))
        conn.cursor().execute(sql.COPY_CSV)
        conn.cursor().execute(sql.UPDATE_TMP_TABLE)
        conn.cursor().execute(sql.UNLOAD_CSV.format(
            max_file_size=config.MAX_FILE_SIZE))
        conn.cursor().execute(sql.GET_CSV.format(file_name=output_dir))
    except SnowflakeError:
        conn.rollback()
        log.exception('Failed to process Snowflake data.')
        sys.exit(1)
    finally:
        conn.close()


def process_snowflake_for_fingerprinting(file_name, output_dir, conn=None):
    """Import CSV data into snowflake and get the corresponding TUIDs.
    Use fingerprinting tables instead of registry for the update step.

    The steps are:
    1. Import the CSV into Snowflake tmp table
    2. Update tmp table with tuid and internal conflict information
    from facts.prod.registry
    3. Export an updated tmp table to CSV
    4. Download the exported CSV file

    Args:
        file_name (str): csv file to process
        output_dir (str): output directory to save the processed data

    """
    if conn is None:
        conn = snowflake.get_snowflake_connection()
    try:
        log.info('Creating snowflake table.')
        conn.cursor().execute(sql_fp.CREATE_TMP_TABLE)
        log.info('Putting CSV into snowflake.')
        conn.cursor().execute(sql_fp.PUT_CSV.format(file_name=file_name))
        log.info('Copy data from CSV into snowflake table.')
        conn.cursor().execute(sql_fp.COPY_CSV)
        log.info('Updating snowflake tmp table with tuid and conflict.')
        conn.cursor().execute(sql_fp.UPDATE_TMP_TABLE)
        log.info('Unloading CSV file from Snowflake.')
        conn.cursor().execute(sql_fp.UNLOAD_CSV.format(
            max_file_size=config.MAX_FILE_SIZE))
        log.info(f'Writing csv to {output_dir}.')
        conn.cursor().execute(sql_fp.GET_CSV.format(file_name=output_dir))
    except SnowflakeError:
        conn.rollback()
        log.exception('Failed to process Snowflake data.')
        sys.exit(1)
    finally:
        conn.close()


def import_csv(file_name):
    """Import CSV file with updated information into MySQL stmt-db.

    Args:
        file_name (str): csv file name with imported data

    """
    conn = mysql.get_connection()
    # TODO: Remove this once fingerprinting data is live
    if config.USE_MRR:
        fieldnames = [
            const.ISRC,
            const.TERRITORY,
            const.TUID,
            const.INTERNAL_CONFLICT,
        ]
    else:
        fieldnames = [
            const.SERVICE,
            const.DATE_COLUMN_NAME,
            const.ISRC,
            const.TERRITORY,
            const.TUID,
            const.INTERNAL_CONFLICT,
            const.RULES_SUMMARY,
        ]
    batch_size = (
        config.BATCH_SIZE if config.USE_MRR else config.BULK_BATCH_SIZE)
    try:
        with gzip.open(file_name, 'rt') as f:
            reader = csv.DictReader(f, fieldnames=fieldnames)
            records = util.grouper(batch_size, reader)
            for chunk in records:
                statement_db.update_report(conn, chunk)
                log.info(f'Updated {batch_size} records.')
    except err.MySQLError:
        log.exception('Failed to import data into MySQL')
        sys.exit(1)
    finally:
        conn.close()


def start():
    """Start processing."""
    with tempfile.TemporaryDirectory(dir=os.getcwd()) as tmpdirname:
        # TODO: Move this back into process_snowflake once pw auth is removed
        conn = snowflake.get_snowflake_connection()
        export_filename = f'{tmpdirname}/{const.EXPORT_FILENAME}'
        log.info(f'Created temporary directory {tmpdirname}')
        log.info('Exporting CSV with statement data...')
        rows_exported = export_csv(export_filename)
        log.info(f'{rows_exported} written to CSV.')
        if rows_exported == 0:
            log.info('No data to process.')
            sys.exit(0)
        log.info('Importing and processing the data in Snowflake')
        # TODO: Remove this once fingerprinting data is live
        if config.USE_MRR:
            process_snowflake(export_filename, tmpdirname, conn)
        else:
            process_snowflake_for_fingerprinting(export_filename, tmpdirname,
                                                 conn)
        log.info('Snowflake processing finished.')
        log.info('Importing the data to MySQL.')
        if config.USE_MRR:
            import_csv(f'{tmpdirname}/data')
        else:
            for csv_file in sorted(glob.glob(f'{tmpdirname}/*.gz')):
                import_csv(csv_file)
        log.info('MySQL data import finished.')


if __name__ == '__main__':
    start_time = datetime.utcnow()
    mrr = 'using mrr' if config.USE_MRR else 'using fingerprinting'
    log.info('Job started  '+mrr)
    start()
    elapsed_time = datetime.utcnow() - start_time
    log.info(f'Job finished in {elapsed_time}.')