import asyncio
import io
import logging
import os
import json
import sys
from timeit import default_timer as timer

import asyncpg
import boto3

# Connection pool is shared between concurrent warm Lambda runs;
# can't perform await outside async function, using lazy load below
POOL: asyncpg.Pool = None

CONCURRENCY = 5
MAX_POOL_SIZE = 100
MAX_ROWS_IN_BATCH = 1000


class FanAttributeCSVToDatabaseSaver:
    def __init__(self, schema, table, collection_id, file_path, set_finished_status=True):
        self.schema = schema
        self.table = table
        self.collection_id = collection_id
        self.file_path = file_path
        self.rows_inserted_total = 0
        self.set_finished_status = set_finished_status
        self._semaphore = asyncio.Semaphore(CONCURRENCY)

    def _get_secret(self, secret_name):
        """
        This needs AWS default config (or profile named "fansifter")
        """
        region_name = "eu-west-1"
        session = boto3.session.Session()
        client = session.client(service_name="secretsmanager", region_name=region_name)
        response = client.get_secret_value(SecretId=secret_name)["SecretString"]
        return response

    def _get_connection_string(self):
        profile = os.getenv("PROFILE")
        if profile == "devel":
            secret_name = "fansifter-rds"
        elif profile == "test":
            secret_name = "fansifter-rds-test"
        elif profile == "live":
            secret_name = "fansifter-rds-live"
        else:
            raise ValueError(f"Unknown profile: {profile}")
        secret = json.loads(self._get_secret(secret_name))
        return "postgres://{username}:{password}@{host}:{port}/{dbname}".format(**secret)

    async def _connect_to_database(self):
        global POOL
        if POOL is None:  # lazy initialization of PG Connection Pool
            POOL = await asyncpg.create_pool(self._get_connection_string(), min_size=5, max_size=MAX_POOL_SIZE)

    async def _execute_raw_query(self, query, fetch=False, exit_on_failure=True):
        async with POOL.acquire() as connection:
            try:
                if fetch:
                    return await connection.fetch(query)
                else:
                    return await connection.execute(query)
            except Exception as e:
                logging.error(e, exc_info=e)
                if exit_on_failure:
                    sys.exit(1)

    async def _delete_existing_data(self):
        [[count]] = await self._execute_raw_query(
            f"""
                SELECT COUNT(1) FROM {self.schema}.{self.table}
                WHERE collection_id = {self.collection_id};
        """,
            fetch=True,
        )
        if count > 0:
            await self._execute_raw_query(
                f"""
                    DELETE FROM {self.schema}.{self.table}
                    WHERE collection_id = {self.collection_id};
                """
            )

    async def _save_bytes_buffer_to_db(self, buffer):
        try:
            async with POOL.acquire() as connection:  # taking a connection from pool
                buffer.seek(0)  # move position to beginning of file before reading
                await connection.copy_to_table(
                    table_name=self.table,
                    source=buffer,
                    schema_name=self.schema,
                    format="csv",
                    delimiter="\t",
                    header="False",
                )
        except Exception as e:
            logging.error(e, exc_info=e)
            sys.exit(1)
        finally:
            self._semaphore.release()  # allowing next insertion of a chunk to start

    async def _read_file_and_save_in_chunks(self):
        with open(self.file_path, "r") as f:
            await self._delete_existing_data()
            tasks = []  # asyncio Futures
            buffer = io.BytesIO()
            rows_in_buffer = 0
            for line in f:
                buffer.write(line.encode("utf-8"))  # storing lines as bytes
                rows_in_buffer += 1
                self.rows_inserted_total += 1
                if rows_in_buffer == MAX_ROWS_IN_BATCH:
                    await self._semaphore.acquire()  # stop reading while all %CONCURRENCY% insertions are in progress
                    tasks.append(
                        asyncio.ensure_future(
                            # schedule concurrent async insertion
                            self._save_bytes_buffer_to_db(buffer)
                        )
                    )
                    buffer = io.BytesIO()  # recreating buffer to start filling a new one
                    rows_in_buffer = 0
            if rows_in_buffer > 0:  # insert rows from last chunk
                tasks.append(asyncio.ensure_future(self._save_bytes_buffer_to_db(buffer)))

            await asyncio.gather(*tasks)  # waiting for all insertions to finish

    async def _update_collection_status(self, status):
        return await self._execute_raw_query(
            f"""
                UPDATE {self.schema}.collection
                SET status = '{status}'
                WHERE id = {self.collection_id};
            """
        )

    def _remove_file(self):
        try:
            if os.path.exists(self.file_path):
                os.remove(self.file_path)
        except Exception as e:
            logging.error(e, exc_info=e)

    async def save_to_database(self):
        await self._connect_to_database()

        start = timer()

        await self._read_file_and_save_in_chunks()
        if self.set_finished_status:
            await self._update_collection_status("finished")
        self._remove_file()

        end = timer()
        print(
            f"Inserting {self.rows_inserted_total} rows of collection {self.collection_id} into [{os.getenv('PROFILE')}]{self.schema}.{self.table} took {end - start:.2f} seconds."
        )


async def main(event, context):
    saver = FanAttributeCSVToDatabaseSaver(
        event["schema"], event["table"], event["collection_id"], event["file_path"], event["set_finished_status"]
    )
    await saver.save_to_database()
    return {"statusCode": 200, "body": json.dumps({"success": True, "rows_inserted": saver.rows_inserted_total})}


def lambda_handler(event, context):
    """
    :param event: processing parameters
    :param event.schema: database schema
    :param event.table: table to insert data into
    :param event.collection_id: collection id
    :param event.file_path full path to csv file (EFS volume is mounted into /mnt/efs)
    :param event.set_finished_status whether to set collection status as "finished"
    :param context: ignored
    """
    loop = asyncio.get_event_loop()
    return loop.run_until_complete(main(event, context))