#
# Copyright (c) 2022 Airbyte, Inc., all rights reserved.
#

import json

import boto3
from airbyte_cdk.destinations import Destination
from botocore.exceptions import ClientError
from retrying import retry

from .config_reader import AuthMode, ConnectorConfig


class AwsHandler:
    COLUMNS_MAPPING = {"number": "float", "string": "string", "integer": "int"}

    def __init__(self, connector_config, destination: Destination):
        self._connector_config: ConnectorConfig = connector_config
        self._destination: Destination = destination
        self._bucket_name = connector_config.bucket_name
        self.logger = self._destination.logger

        self.create_session()
        self.s3_client = self.session.client("s3", region_name=connector_config.region)
        self.glue_client = self.session.client("glue")
        self.lf_client = self.session.client("lakeformation")

    @retry(stop_max_attempt_number=10, wait_random_min=1000, wait_random_max=2000)
    def create_session(self):
        if self._connector_config.credentials_type == AuthMode.IAM_USER.value:
            self._session = boto3.Session(
                aws_access_key_id=self._connector_config.aws_access_key,
                aws_secret_access_key=self._connector_config.aws_secret_key,
                region_name=self._connector_config.region,
            )
        elif self._connector_config.credentials_type == AuthMode.IAM_ROLE.value:
            client = boto3.client("sts")
            role = client.assume_role(
                RoleArn=self._connector_config.role_arn,
                RoleSessionName="airbyte-destination-aws-datalake",
            )
            creds = role.get("Credentials", {})
            self._session = boto3.Session(
                aws_access_key_id=creds.get("AccessKeyId"),
                aws_secret_access_key=creds.get("SecretAccessKey"),
                aws_session_token=creds.get("SessionToken"),
                region_name=self._connector_config.region,
            )
        else:
            raise Exception("Session wasn't created")

    @property
    def session(self) -> boto3.Session:
        return self._session

    @retry(stop_max_attempt_number=10, wait_random_min=2000, wait_random_max=3000)
    def head_bucket(self):
        self.s3_client.head_bucket(Bucket=self._bucket_name)

    @retry(stop_max_attempt_number=10, wait_random_min=2000, wait_random_max=3000)
    def head_object(self, object_key):
        return self.s3_client.head_object(Bucket=self._bucket_name, Key=object_key)

    @retry(stop_max_attempt_number=10, wait_random_min=2000, wait_random_max=3000)
    def put_object(self, object_key, body):
        self.s3_client.put_object(Bucket=self._bucket_name, Key=object_key, Body="\n".join(body))

    @staticmethod
    def batch_iterate(iterable, n=1):
        size = len(iterable)
        for ndx in range(0, size, n):
            yield iterable[ndx : min(ndx + n, size)]

    def get_table(self, txid, database_name: str, table_name: str, location: str):
        table = None
        try:
            table = self.glue_client.get_table(DatabaseName=database_name, Name=table_name, TransactionId=txid)
        except ClientError as e:
            if e.response["Error"]["Code"] == "EntityNotFoundException":
                table_input = {
                    "Name": table_name,
                    "TableType": "GOVERNED",
                    "StorageDescriptor": {
                        "Location": location,
                        "InputFormat": "org.apache.hadoop.mapred.TextInputFormat",
                        "OutputFormat": "org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat",
                        "SerdeInfo": {"SerializationLibrary": "org.openx.data.jsonserde.JsonSerDe", "Parameters": {"paths": ","}},
                    },
                    "PartitionKeys": [],
                    "Parameters": {"classification": "json", "lakeformation.aso.status": "true"},
                }
                self.glue_client.create_table(DatabaseName=database_name, TableInput=table_input, TransactionId=txid)
                table = self.glue_client.get_table(DatabaseName=database_name, Name=table_name, TransactionId=txid)
            else:
                err = e.response["Error"]["Code"]
                self.logger.error(f"An error occurred: {err}")
                raise

        if table:
            return table
        else:
            return None

    def update_table(self, database, table_info, transaction_id):
        self.glue_client.update_table(DatabaseName=database, TableInput=table_info, TransactionId=transaction_id)

    def preprocess_type(self, property_type):
        if type(property_type) is list:
            not_null_types = list(filter(lambda t: t != "null", property_type))
            if len(not_null_types) > 2:
                return "string"
            else:
                return not_null_types[0]
        else:
            return property_type

    def cast_to_athena(self, str_type):
        preprocessed_type = self.preprocess_type(str_type)
        return self.COLUMNS_MAPPING.get(preprocessed_type, preprocessed_type)

    def generate_athena_schema(self, schema):
        columns = []
        for (k, v) in schema.items():
            athena_type = self.cast_to_athena(v["type"])
            if athena_type == "object":
                properties = v["properties"]
                type_str = ",".join([f"{k1}:{self.cast_to_athena(v1['type'])}" for (k1, v1) in properties.items()])
                columns.append({"Name": k, "Type": f"struct<{type_str}>"})
            else:
                columns.append({"Name": k, "Type": athena_type})
        return columns

    def update_table_schema(self, txid, database, table, schema):
        table_info = table["Table"]
        table_info_keys = list(table_info.keys())
        for k in table_info_keys:
            if k not in [
                "Name",
                "Description",
                "Owner",
                "LastAccessTime",
                "LastAnalyzedTime",
                "Retention",
                "StorageDescriptor",
                "PartitionKeys",
                "ViewOriginalText",
                "ViewExpandedText",
                "TableType",
                "Parameters",
                "TargetTable",
                "IsRowFilteringEnabled",
            ]:
                table_info.pop(k)

        self.logger.debug("Schema = " + repr(schema))

        columns = self.generate_athena_schema(schema)
        if "StorageDescriptor" in table_info:
            table_info["StorageDescriptor"]["Columns"] = columns
        else:
            table_info["StorageDescriptor"] = {"Columns": columns}
        self.update_table(database, table_info, txid)
        self.glue_client.update_table(DatabaseName=database, TableInput=table_info, TransactionId=txid)

    def get_all_table_objects(self, txid, database, table):
        table_objects = []

        try:
            res = self.lf_client.get_table_objects(DatabaseName=database, TableName=table, TransactionId=txid)
        except ClientError as e:
            if e.response["Error"]["Code"] == "EntityNotFoundException":
                return []
            else:
                err = e.response["Error"]["Code"]
                self.logger.error(f"Could not get table objects due to error: {err}")
                raise

        while True:
            next_token = res.get("NextToken", None)
            partition_objects = res.get("Objects")
            table_objects.extend([p["Objects"] for p in partition_objects])
            if next_token:
                res = self.lf_client.get_table_objects(
                    DatabaseName=database,
                    TableName=table,
                    TransactionId=txid,
                    NextToken=next_token,
                )
            else:
                break
        flat_list = [item for sublist in table_objects for item in sublist]
        return flat_list

    def purge_table(self, txid, database, table):
        self.logger.debug(f"Going to purge table {table}")
        write_ops = []
        all_objects = self.get_all_table_objects(txid, database, table)
        write_ops.extend([{"DeleteObject": {"Uri": o["Uri"]}} for o in all_objects])
        if len(write_ops) > 0:
            self.logger.debug(f"{len(write_ops)} objects to purge")
            for batch in self.batch_iterate(write_ops, 99):
                self.logger.debug("Purging batch")
                try:
                    self.lf_client.update_table_objects(
                        TransactionId=txid,
                        DatabaseName=database,
                        TableName=table,
                        WriteOperations=batch,
                    )
                except ClientError as e:
                    self.logger.error(f"Could not delete object due to exception {repr(e)}")
                    raise
        else:
            self.logger.debug("Table was empty, nothing to purge.")

    def update_governed_table(self, txid, database, table, bucket, object_key, etag, size):
        self.logger.debug(f"Updating governed table {database}:{table}")
        write_ops = [
            {
                "AddObject": {
                    "Uri": f"s3://{bucket}/{object_key}",
                    "ETag": etag,
                    "Size": size,
                }
            }
        ]

        self.lf_client.update_table_objects(
            TransactionId=txid,
            DatabaseName=database,
            TableName=table,
            WriteOperations=write_ops,
        )


class LakeformationTransaction:
    def __init__(self, aws_handler: AwsHandler):
        self._aws_handler = aws_handler
        self._transaction = None
        self._logger = aws_handler.logger

    @property
    def txid(self):
        return self._transaction["TransactionId"]

    def cancel_transaction(self):
        self._logger.debug("Canceling Lakeformation Transaction")
        self._aws_handler.lf_client.cancel_transaction(TransactionId=self.txid)

    def commit_transaction(self):
        self._logger.debug(f"Commiting Lakeformation Transaction {self.txid}")
        self._aws_handler.lf_client.commit_transaction(TransactionId=self.txid)

    def extend_transaction(self):
        self._logger.debug("Extending Lakeformation Transaction")
        self._aws_handler.lf_client.extend_transaction(TransactionId=self.txid)

    def describe_transaction(self):
        return self._aws_handler.lf_client.describe_transaction(TransactionId=self.txid)

    def __enter__(self, transaction_type="READ_AND_WRITE"):
        self._logger.debug("Starting Lakeformation Transaction")
        self._transaction = self._aws_handler.lf_client.start_transaction(TransactionType=transaction_type)
        self._logger.debug(f"Transaction id = {self.txid}")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._logger.debug("Exiting LakeformationTransaction context manager")
        tx_desc = self.describe_transaction()
        self._logger.debug(json.dumps(tx_desc, default=str))

        if exc_type:
            self._logger.error("Exiting LakeformationTransaction context manager due to an exception")
            self._logger.error(repr(exc_type))
            self._logger.error(repr(exc_val))
            self.cancel_transaction()
            self._transaction = None
        else:
            self._logger.debug("Exiting LakeformationTransaction context manager due to reaching end of with block")
            try:
                self.commit_transaction()
                self._transaction = None
            except Exception as e:
                self.cancel_transaction()
                self._logger.error(f"Could not commit the transaction id = {self.txid} because of :\n{repr(e)}")
                self._transaction = None
                raise (e)
