import logging
import operator
from io import BytesIO, StringIO
from typing import Optional

import boto3
import charset_normalizer
import pandas as pd
import psycopg2
import sqlalchemy
from numpy import ndarray
from pandas.io import sql
from psycopg2.extras import execute_values

from service.db import engine

logger = logging.getLogger(__name__)


def pd_read_sql(query_sql, **kwargs) -> pd.DataFrame:
    attempts = 10
    # workaround for sqlalchemy.exc.OperationalError: (psycopg2.OperationalError) SSL connection has been closed
    # unexpectedly
    while attempts > 0:
        try:
            attempts -= 1
            df = pd.read_sql(sql=query_sql, con=engine, **kwargs)
            return df
        except sqlalchemy.exc.OperationalError as e:
            logger.exception(e)
        except psycopg2.OperationalError as e:
            logger.exception(e)
        except Exception as e:
            logger.exception(e)
            raise e


def run_query(
    query_sql, *args, return_type=None, update_values=False, fetch=True, **kwargs
):
    """
    Runs query against AWS RDS PostgresSQL instance (version >10).
    Creates connection, runs query, ommmits, terminates connection. If there are rows, returns them.
    Avoid creating connection for each individual INSERT, or refactor this function.
    Need to explicitly grant all privileges on any new schemas and tables when not superuser.
    """
    attempts = 10
    # workaround for sqlalchemy.exc.OperationalError: (psycopg2.OperationalError) SSL connection has been closed
    # unexpectedly
    while attempts > 0:
        try:
            attempts -= 1
            connection = None
            for attempt in range(10):
                if connection is not None:
                    logger.warning(f"Retry connection, attempt {attempt}")
                    try:
                        connection.cancel()
                        connection.reset()
                        connection.close()
                    except Exception as e:
                        logger.exception(e)
                connection = engine.raw_connection()
                if not (
                    connection.connection.closed or connection.connection.info.status
                ):
                    # info.status comes from C library, 0 OK 1 BAD
                    break
            cursor = connection.cursor()

            # Using parameterized query to avoid sql injection and other similar issues
            # E.g. it fixes this, which before caused syntax error: "text": "Economy of the People\'s <-- the apostrophe
            # If there are args, it comes from insert and we insert a tuple. This seems easy to break, so add test cases
            # print(cursor.mogrify(query_sql, (args,)))  # for when we want to see the sql
            records = None

            if update_values and args and isinstance(args[0], list):
                # TODO: probably works smoother with pandas?
                execute_values(cursor, query_sql, args[0])
            else:
                if return_type is None:
                    if args and isinstance(args[0], dict):
                        cursor.execute(query_sql, args[0])
                    elif args and isinstance(args[0], (list, ndarray)):
                        # cursor.executemany(query_sql, args[0])
                        records = execute_values(
                            cursor,
                            query_sql,
                            args[0],
                            page_size=len(args[0]),
                            fetch=fetch,
                        )
                        return_type = "toomany"
                    elif isinstance(query_sql, list):
                        for ql in query_sql:
                            cursor.execute(ql, *args, **kwargs)
                    else:
                        cursor.execute(query_sql, *args, **kwargs)

                try:
                    # See if we want to return dataframe instead of list. Consider only using dataframes.
                    if return_type == "df":
                        if args and isinstance(args[0], dict):
                            records = sql.read_sql(
                                query_sql, connection, params=args[0]
                            )
                        else:
                            records = sql.read_sql(query_sql, connection)
                    elif return_type == "dict":
                        records = sql.read_sql(query_sql, connection).to_dict(
                            orient="records"
                        )
                    elif return_type == "toomany":
                        pass
                    elif fetch:
                        records = cursor.fetchall()
                except psycopg2.ProgrammingError:
                    pass
                except Exception as e:
                    # TODO! this needs to be handled more gracefully (instead pass)
                    #  We want don't want to fetch when we insert
                    logger.exception(f"FAILED run_query {query_sql}", exc_info=e)

            connection.commit()
            connection.close()
            return records

        except sqlalchemy.exc.OperationalError as e:
            logger.exception(e)
        except psycopg2.OperationalError as e:
            logger.exception(e)
        except Exception as e:
            logger.exception(e)
            raise e


def s3_to_pandas(
    bucket: str, key: str, header: Optional[str] = None, file_format: str = "csv"
) -> pd.DataFrame:
    """
    Downloads S3 file and returns it in a pandas DataFrame.
    If you want to get header from file, pass header="infer".
    """
    logger.info(f"Attempting to read {key} file from {bucket} bucket...")
    client = boto3.client("s3")
    body_bytes = client.get_object(Bucket=bucket, Key=key)["Body"].read()

    if file_format == "csv":
        best_encoding_match = charset_normalizer.from_bytes(body_bytes).best()
        if best_encoding_match is not None:
            # we set encoding as an S3 object tag
            # to be used when reading the file next time
            encoding = best_encoding_match.encoding
            client.put_object_tagging(
                Bucket=bucket,
                Key=key,
                Tagging={"TagSet": [{"Key": "encoding", "Value": encoding}]},
            )
            normalized_bytes = best_encoding_match.output()
        else:
            # If the encoding was not guessed by charset_normalizer,
            # let's hope pandas will guess it.
            normalized_bytes = body_bytes

        # Although csv-sniffer in C and Python should do that automatically, they can
        # fail sometimes, We count the occurances of ; , \t and guess
        # which one may be a delimiter.
        first_256_chars = str(normalized_bytes[:256])
        char_stats = {
            ";": first_256_chars.count(";"),
            ",": first_256_chars.count(","),
            "\t": first_256_chars.count("\t"),
        }
        delimiter = max(char_stats.items(), key=operator.itemgetter(1))[0]

        return pd.read_csv(
            BytesIO(normalized_bytes), sep=delimiter, header=header, dtype=str
        )
    else:
        # TODO add global support for xls/xlsx with pd.read_excel
        raise ValueError(f"Unsupported file format: {file_format}")


def pandas_to_s3(bucket, df, key, index=False):
    """Creates a CSV from pandas DataFrame and uploads it to S3"""
    csv_buffer = StringIO()
    df.to_csv(csv_buffer, index=index)
    s3_resource = boto3.resource("s3")
    s3_resource.Object(bucket, key).put(Body=csv_buffer.getvalue())


def df_to_db(
    df: pd.DataFrame, schema: str, table_name: str, if_exists="append", index_label=None
):
    """
    Using df.to_sql to insert the df in to table
    :param df: source_df
    :param schema: schema name
    :param table_name: table name
    :param if_exists: if exists parameter default: 'append'
    :param index_label: index label
    :return:
    """
    df.to_sql(
        name=table_name,
        con=engine,
        schema=schema,
        if_exists=if_exists,
        chunksize=16000,
        index=False,
        index_label=index_label,
    )
    return