"""Close Batch Lambda index module."""

import base64
import json
import os
import tempfile
import time

from sentry_sdk import capture_exception
from vector_utils.connections import connection_info
from vector_utils.connections import exceptions
from vector_utils.connections.transporter import Transporter
from vector_utils import queries
from vector_utils import utils
from vector_utils.datadog import metrics
from vector_utils.utils import parsers
from vector_utils.utils import write_file_to_disk
from src import config
from src.config import vector_secrets_manager_client
from src.logger import get_current_logger
from src.manifestor import Manifestor


replacements = {
    'Year': time.strftime('%Y'),
    'year': time.strftime('%y'),
    'month': time.strftime('%m'),
    'day': time.strftime('%d'),
    'hour': time.strftime('%H'),
    'minute': time.strftime('%M'),
    'second': time.strftime('%S')
}
correlation_id = None


def handler(event, context):
    """Handle event.

    Args:
        event (optional): AWS Lambda event dependent structure with metadata.
        context (LambdaContext): AWS Lambda context.
    """
    correlation_id = context.aws_request_id
    logger = get_current_logger(correlation_id)

    try:
        metrics.call_datadog_with_metric(
            metric=f'{config.DATADOG_SCRIPT_NAME}.{config.DATADOG_ATTEMPT}',
            tags=['environment:{}'.format(config.ENVIRONMENT)],
            api_key=config.DATADOG_API_KEY,
            app_key=config.DATADOG_APP_KEY
        )

        messages = event['Records']

        if not messages:
            logger.info('No messages to process')
            return

        message = messages[0]
        message_body = json.loads(message['body'])
        batch_id = message_body['batch_id']

        if not batch_id:
            raise ValueError('Batch ID must be greater than 0.')

        logger.info(f'Batch: {batch_id} update status to in-flight')
        result = queries.update_batch_status(
            'queued_for_closing', 'in-flight', batch_id,
            conn_info=config.DD_MYSQL_CONN_INFO)

        if not result:
            logger.error(
                f'Batch id: {batch_id} Failed to update status from '
                'queued_for_closing to in-flight')
            return
        try:
            close_batch(batch_id)
        except Exception as e:
            queries.update_batch_status(
                'in-flight', 'open', batch_id,
                conn_info=config.DD_MYSQL_CONN_INFO)
            logger.info(f'Batch id: {batch_id} reset to open status due to {e}')
    except Exception as e:
        capture_exception(e)
        logger.error(e)
        raise e


def call_datadog_with_metric(metric, dms_id):
    """Call datadog with metric.

    Args:
        dms_id (str): store id.
        metric (str): custom metric name to send to DataDog.
    """
    if config.DISABLE_PUBLISHING_METRICS:
        return
    metrics.call_datadog_with_metric(
        metric=f'{config.DATADOG_SCRIPT_NAME}.{metric}',
        tags=[
            'environment:{}'.format(config.ENVIRONMENT),
            'dms_id:{}'.format(dms_id)
        ],
        api_key=config.DATADOG_API_KEY,
        app_key=config.DATADOG_APP_KEY
    )


def process_close_batch_files(file_list, batch_info, dms_delivery_spec, secret_data=None):
    """Process a list of files to close batch on.

    Args:
        file_list (list): File list
        batch_info (dict): Batch information
        dms_delivery_spec (dict): DMS delivery spec
        secret_data (dict): Data from Secret manager
    """
    log = get_current_logger(correlation_id)

    if not file_list:
        log.info('No files to close batch on')
        return

    dms_id = batch_info['dms_id']
    order_type = batch_info['order_type']
    batch_id = batch_info['batch_id']

    call_datadog_with_metric(config.DATADOG_TRANSFER_START, dms_id)
    log.info(f'Closing batch for DMS: {dms_id} for batch_id: {batch_id}')
    connection_data = secret_data

    if 'private_key' in connection_data:
        connection_data['priv_key'] = f'{tempfile.gettempdir()}/{dms_id}_{order_type}_key'

        write_file_to_disk(
            tempfile.gettempdir(),
            connection_data['priv_key'],
            base64.b64decode(connection_data['private_key']).decode())

    if (
        connection_data.get('connection_type') == 'sftp'
        and str(dms_id) in config.STORES_WITH_DISABLED_ALGORITHMS
    ):
        log.info('Set sftp_disabled_algorithms for: {}'.format(dms_id))
        connection_data['sftp_disabled_algorithms'] = {
            'pubkeys': ['rsa-sha2-256', 'rsa-sha2-512']}

    if connection_data.get('connection_type') == 'gcs' and 'gcs_config_file' in connection_data:
        gcs_config = connection_data['gcs_config_file']
        if isinstance(gcs_config, dict) and 'private_key' in gcs_config:
            try:
                decoded_key = base64.b64decode(gcs_config['private_key']).decode()
                gcs_config['private_key'] = decoded_key
                connection_data['gcs_config_file'] = gcs_config
            except Exception as e:
                log.error(f'Failed to decode GCS private key for {dms_id}: {e}')
                raise ValueError(f'Invalid base64 private key for GCS config (dms_id={dms_id}): {e}')

    conn_obj = connection_info.ConnectionInfo(connection_data, order_type)

    log.info('Creating transporter with conn obj: {}'.format(conn_obj))

    try:
        transporter = Transporter(conn_obj)
    except (exceptions.ConnectionTimeout, exceptions.SSHException) as err:
        log.warning(f'Error {err} for DMS {dms_id}')
        call_datadog_with_metric(config.DATADOG_TRANSFER_FAIL, dms_id)
        raise err

    check_delivered_files = True \
        if dms_delivery_spec['confirm_manifest_delivery'] == 'Y' \
        else False

    try:
        transporter.transfer_files(
            file_list, transfer_mode='upload',
            batch_file_id=batch_info['dms_id'],
            check_delivered_files=check_delivered_files)
    except (exceptions.ConnectionTimeout, exceptions.SSHException, exceptions.TransferException) as err:
        error_message = f'Attempted to upload to server {dms_delivery_spec["domain_name"]} '
        f'for DMS {dms_id} and failed after 5 tries.'
        log.warning(error_message)
        call_datadog_with_metric(config.DATADOG_TRANSFER_FAIL, dms_id)
        raise err
    else:
        log.info(f'Dropped delivery.complete file for DMS: {dms_id} and batch_id: {batch_id}')
        call_datadog_with_metric(config.DATADOG_TRANSFER_SUCCESS, dms_id)

        queries.close_delivery_batch(
            batch_info['batch_id'], conn_info=config.DD_MYSQL_CONN_INFO)
    finally:
        try:
            transporter.close_connection()
        except Exception:
            pass


def close_batch(batch_id):
    """Entry point to check if we should process a close batch.

    Args:
        batch_id (int): The batch ID
    Returns:
    """
    log = get_current_logger(correlation_id)
    log.info(f'get_batch_info for: {batch_id}')
    batch_info = get_batch_info(batch_id)
    if not batch_info:
        return

    dms_id = batch_info['dms_id']
    order_type = batch_info['order_type']

    try:
        dms_delivery_spec = queries.get_dms_delivery_spec(
            dms_id, order_type,
            conn_info=config.DD_MYSQL_CONN_INFO)

        secrets_path = f'connection_info/{dms_id}/{order_type}'
        secret_data = vector_secrets_manager_client.get_cred(secrets_path)
        if not secret_data:
            raise Exception(f'Connection info missing at: {secrets_path}')
        dms_delivery_spec['remote_initial_dir'] = secret_data['remote_initial_dir']

        file_list = []
        replacements['batch_folder'] = batch_info['remote_folder']

        if dms_delivery_spec['manifest_filename']:
            manifest_filename = parsers.process_string(
                dms_delivery_spec['manifest_filename'],
                replacements)

            manifestor = Manifestor(
                dms_delivery_spec['manifest_format'],
                dms_delivery_spec['remote_initial_dir'],
                batch_info)

            generate_local_file(
                file_list,
                dms_delivery_spec['remote_initial_dir'],
                batch_info,
                manifest_filename,
                manifestor.generate_manifest())

        if dms_delivery_spec['delivery_complete_file']:
            delivery_complete_file = parsers.process_string(
                dms_delivery_spec['delivery_complete_file'],
                replacements)

            generate_local_file(
                file_list,
                dms_delivery_spec['remote_initial_dir'],
                batch_info,
                delivery_complete_file,
                '')

        process_close_batch_files(file_list, batch_info, dms_delivery_spec, secret_data)
    except Exception as e:
        call_datadog_with_metric(config.DATADOG_TRANSFER_FAIL, dms_id)
        raise e


def get_batch_info(batch_id):
    """Get batch information.

    Args:
        batch_id (int): Batch ID

    Returns:
        dict or None: Batch information
    """
    log = get_current_logger(correlation_id)
    results = queries.get_batch_info(
        batch_id, conn_info=config.DD_MYSQL_CONN_INFO)

    if not results:
        queries.close_delivery_batch(
            batch_id, conn_info=config.DD_MYSQL_CONN_INFO)
        log.info(f'No batch info found with delivered status in batch: {batch_id}')
        return None

    batch_info = {}
    for row in results:
        meta_update = True if row['meta_update'] == 'Y' else False
        job_detail = {
            'eqd_id': row['encoding_queue_detail_id'],
            'upc': row['upc'],
            'meta_update': meta_update,
        }

        if not batch_info:
            batch_info = {
                'batch_id': batch_id,
                'dms_id': row['dms_master_master_id'],
                'order_type': row['order_type'],
                'remote_folder': row['remote_folder'],
                'delivered_upcs': [job_detail]
            }
        else:
            batch_info['delivered_upcs'].append(job_detail)
    return batch_info


def generate_local_file(
        file_list, remote_initial_dir, batch_info, file_name, file_data):
    """Generate local file.

    Args:
        file_list (list): List of local and remote file pairs
        remote_initial_dir (str): Remote initial folder
        batch_info (dict): Batch info
        file_name (str): Local file name
        file_data (str): File content

    Returns:
        list
    """
    local_file = os.path.basename(file_name)
    temp_path = tempfile.gettempdir()
    temp_file = '{dms_id}_{batch_id}_{filename}'.format(
        dms_id=batch_info['dms_id'],
        batch_id=batch_info['batch_id'],
        filename=local_file)
    utils.write_file_to_disk(temp_path, temp_file, file_data)

    file_list.append({
        'local': '{}/{}'.format(temp_path, temp_file),
        'remote': '/'.join([remote_initial_dir, file_name])
    })


if __name__ == '__main__':
    handler(None, None)