"""Entrypoint."""

import os

from extract_sales import config
from extract_sales.connectors.logger import logger
from extract_sales.connectors.sentry import init_sentry
from extract_sales.connectors.ssh import get_ssh_conn
from extract_sales.connectors.ssh import run_ssh_cmd
from extract_sales.constants import SALES_TYPE_TO_TABLE_NAME
from extract_sales.templates import SELECT_SALES_BATCH_TEMPLATE

if config.SENTRY_DSN:
    init_sentry(config.SENTRY_DSN)


def run():
    """Extract sales from StatementDB to S3.

    - Connect to the StatementDB EC2 instance via SSH
    - Extract the sales data from the MySQL DB to an outfile
    - Split and compress the outfile
    - Upload the outfile to S3
    - Delete files from the remote filesystem
    """
    try:
        logger.info(f'Starting {config.APPLICATION_NAME} in {config.ENVIRONMENT}')

        (sales_type, batch_id) = _get_run_params()

        logger.info(f'Extracting "{sales_type}" sales for batch "{batch_id}"')

        with get_ssh_conn() as ssh_conn:
            (raw_dir, parts_dir) = _create_dirs(ssh_conn, sales_type, batch_id)
            logger.info('Created directories')

            outfile_path = _generate_outfile(ssh_conn, sales_type, batch_id, raw_dir)
            logger.info('Generated outfile')

            _split_and_compress_outfile(ssh_conn, outfile_path, parts_dir)
            logger.info('Split and compressed outfile')

            _upload_parts(ssh_conn, sales_type, batch_id, parts_dir)
            logger.info('Parts uploaded')

            _delete_files(ssh_conn, raw_dir, parts_dir)
            logger.info('Files deleted')
    except Exception as e:
        logger.error(f'Unexpected error: {e}')
        raise e


def _get_run_params():
    """Get the parameters to run the task from the environment variables.

    The parameters are:
        - SALES_TYPE: The type of sales to extract ("distro" or "nr").
        - BATCH_ID: The ID of the sales batch to extract.

    Returns:
        A tuple containing the `sales_type` and the `batch_id`.
    """
    sales_type = os.environ.get('SALES_TYPE')
    batch_id = os.environ.get('BATCH_ID')

    if not sales_type:
        raise Exception('Missing environment variable: `SALES_TYPE`')

    if sales_type not in ['distro', 'nr']:
        raise Exception('Invalid environment variable: `SALES_TYPE`')

    if not batch_id:
        raise Exception('Missing environment variable: `BATCH_ID`')

    return (sales_type, batch_id)


def _create_dirs(ssh_conn, sales_type, batch_id):
    """Create directories on the remote filesystem.

    This creates a `raw` directory to write the outfile to
    and a `parts` directory to write the outfile's parts once
    it's split and compressed.

    Args:
        ssh_conn (fabric.Connection): The SSH connection to StatementDB.
        sales_type (str): The type of sales to extract ("distro" or "nr").
        batch_id (str): The ID of the sales batch to extract.

    Returns:
        A tuple containing the paths to the created directories.
    """
    raw_dir = f'{config.OUTFILE_DIR}/{sales_type}/{batch_id}/raw'
    parts_dir = f'{config.OUTFILE_DIR}/{sales_type}/{batch_id}/parts'
    cmd = f'install -dv -m 0775 -o {config.SERVICE_USER} -g {config.SERVICE_GROUP}'

    run_ssh_cmd(ssh_conn, f'{cmd} {raw_dir}')
    run_ssh_cmd(ssh_conn, f'{cmd} {parts_dir}')

    return (raw_dir, parts_dir)


def _generate_outfile(ssh_conn, sales_type, batch_id, raw_dir):
    """Generate the outfile on the remote host.

    This executes a query against the MySQL database to extract the sales data
    into a text file.

    Args:
        ssh_conn (fabric.Connection): The SSH connection to StatementDB.
        sales_type (str): The type of sales to extract ("distro" or "nr").
        batch_id (str): The ID of the sales batch to extract.
        raw_dir (str): The path to the directory to write the outfile to.

    Returns:
        The path to the outfile.
    """
    outfile_path = f'{raw_dir}/outfile.txt'
    table_name = SALES_TYPE_TO_TABLE_NAME[sales_type]

    query = SELECT_SALES_BATCH_TEMPLATE.render(
        table_name=table_name,
        batch_id=batch_id
    )

    cmd_args = f'-u {config.MYSQL_DB_USER} -p{config.MYSQL_DB_PASS}'

    # NOTE: Using `--quick` to print each row as it is received to not overload
    # the server by keeping the whole result set in memory
    cmd_args += ' --quick'

    # NOTE: Using `-N` to not include column names in the outfile for easier
    # loading into Snowflake
    cmd_args += ' -N'

    cmd = f'mysql -h {config.SSH_HOST} {cmd_args} -e "{query}" > {outfile_path}'

    run_ssh_cmd(ssh_conn, cmd)

    return outfile_path


def _split_and_compress_outfile(ssh_conn, outfile_path, parts_dir):
    """Split and compress the outfile.

    This executes a `split` command on the remote host using the `--filter`
    argument to compress each part.

    The compression is done with `pigz` which is a parallel implementation of `gzip`.

    Args:
        ssh_conn (fabric.Connection): The SSH connection to StatementDB.
        outfile_path (str): The path to the outfile.
        parts_dir (str): The path to the directory to write the parts to.
    """
    parts_path = f'{parts_dir}/outfile_parts_'

    cmd_args = f'-C {config.SPLIT_SIZE}'
    cmd_args += " --filter='/usr/bin/pigz > $FILE.gz'"

    cmd = f'split {cmd_args} {outfile_path} {parts_path}'

    run_ssh_cmd(ssh_conn, cmd)


def _upload_parts(ssh_conn, sales_type, batch_id, parts_dir):
    """Upload the outfile parts to S3.

    This executes a `s5cmd` command on the remote host to upload the compressed
    parts to S3.

    Args:
        ssh_conn (fabric.Connection): The SSH connection to StatementDB.
        sales_type (str): The type of sales to extract ("distro" or "nr").
        batch_id (str): The ID of the sales batch to extract.
        parts_dir (str): The path to the directory where the parts were written to.
    """
    s3_path = f's3://{config.S3_BUCKET_NAME}/extract-sales/{sales_type}/{batch_id}/'

    cmd = f's5cmd cp {parts_dir} {s3_path}'

    run_ssh_cmd(ssh_conn, cmd)


def _delete_files(ssh_conn, raw_dir, parts_dir):
    """Delete created files from the remote filesystem.

    Args:
        ssh_conn (fabric.Connection): The SSH connection to StatementDB.
        raw_dir (str): The path to the directory where the outfile was written to.
        parts_dir (str): The path to the directory where the parts were written to.
    """
    cmd = f'rm -rf {raw_dir} {parts_dir}'

    run_ssh_cmd(ssh_conn, cmd)


if __name__ == '__main__':
    run()