#!/usr/bin/env python3

# Copyright 2022 Pipin Fitriadi <pipinfitriadi@gmail.com>

# Licensed under the Microsoft Reference Source License (MS-RSL)

# This license governs use of the accompanying software. If you use the
# software, you accept this license. If you do not accept the license, do not
# use the software.

# 1. Definitions

# The terms "reproduce," "reproduction" and "distribution" have the same
# meaning here as under U.S. copyright law.

# "You" means the licensee of the software.

# "Your company" means the company you worked for when you downloaded the
# software.

# "Reference use" means use of the software within your company as a reference,
# in read only form, for the sole purposes of debugging your products,
# maintaining your products, or enhancing the interoperability of your
# products with the software, and specifically excludes the right to
# distribute the software outside of your company.

# "Licensed patents" means any Licensor patent claims which read directly on
# the software as distributed by the Licensor under this license.

# 2. Grant of Rights

# (A) Copyright Grant- Subject to the terms of this license, the Licensor
# grants you a non-transferable, non-exclusive, worldwide, royalty-free
# copyright license to reproduce the software for reference use.

# (B) Patent Grant- Subject to the terms of this license, the Licensor grants
# you a non-transferable, non-exclusive, worldwide, royalty-free patent
# license under licensed patents for reference use.

# 3. Limitations

# (A) No Trademark License- This license does not grant you any rights to use
# the Licensor's name, logo, or trademarks.

# (B) If you begin patent litigation against the Licensor over patents that
# you think may apply to the software (including a cross-claim or counterclaim
# in a lawsuit), your license to the software ends automatically.

# (C) The software is licensed "as-is." You bear the risk of using it. The
# Licensor gives no express warranties, guarantees or conditions. You may have
# additional consumer rights under your local laws which this license cannot
# change. To the extent permitted under your local laws, the Licensor excludes
# the implied warranties of merchantability, fitness for a particular purpose
# and non-infringement.

from copy import copy
import csv
from datetime import date, datetime
from io import StringIO
from json import dumps
import logging
from os import getenv, getcwd
from os.path import isfile, join as path_join
import re
from time import sleep
from typing import Generator, Iterable
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse

from jinja2 import Environment, FileSystemLoader
from jinja2.meta import find_referenced_templates, find_undeclared_variables
from sshtunnel import SSHTunnelForwarder
from sqlalchemy import bindparam, create_engine
from sqlalchemy.engine.base import Engine
from sqlalchemy.engine.default import DefaultDialect
from sqlalchemy.exc import OperationalError, StatementError
from sqlalchemy.orm import scoped_session
from sqlalchemy.orm.query import Query as _Query
from sqlalchemy.orm.session import sessionmaker
from sqlalchemy.orm.scoping import ScopedSessionMixin
from sqlalchemy.sql import text
from sqlalchemy.sql.expression import null
from sqlalchemy.types import (
    Boolean,
    Date,
    DateTime,
    Float,
    Integer,
    JSON,
    NullType,
    String
)

from . import deserialize, json_serializer, serialize


def database_uri(
    db_driver='postgresql+psycopg2',
    db_host='localhost',
    db_port=5432,
    db_name='postgres',
    db_user='',
    db_pass='',
    use_charset_utf8=False
):
    if db_driver == 'mysql' and db_port == 5432:
        db_port = 3306

    uri = (
        (
            f'{ db_driver }://{ db_user }:{ db_pass }@'
            f'{ db_host }:{ db_port }/{ db_name }'
        )
        if db_driver else 'sqlite:///voxrow.db'
    )

    # flask sqlalchemy mysql encoding problems
    # https://stackoverflow.com/questions/26577334/flask-sqlalchemy-mysql-encoding-problems
    # SQLAlchemy + MySQL + UTF-8 support - how?
    # https://groups.google.com/forum/#!topic/pylons-discuss/ol2m46kiSYA
    return (
        uri + '?charset=utf8'
        if db_driver == 'mysql'
        and use_charset_utf8
        else uri
    )


def database_uri_from_env(
    db_driver_env='DB_DRIVER',
    db_host_env='DB_HOST',
    db_port_env='DB_PORT',
    db_name_env='DB_NAME',
    db_user_env='DB_USER',
    db_pass_env='DB_PASS',
    use_charset_utf8_env='USE_CHARSET_UTF8',
):
    return database_uri(
        getenv(db_driver_env, 'postgresql+psycopg2'),
        getenv(db_host_env, 'localhost'),
        int(
            getenv(db_port_env, '5432')
        ),
        getenv(db_name_env, 'postgres'),
        getenv(db_user_env, ''),
        getenv(db_pass_env, ''),
        getenv(use_charset_utf8_env, 'FALSE').upper() in ['TRUE', '1']
    )


# Retry failed sqlalchemy queries
# https://stackoverflow.com/questions/53287215/retry-failed-sqlalchemy-queries/60614707#60614707
class RetryingQuery(_Query):
    __max_retry_count__ = 3

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def __iter__(self):
        attempts = 0

        while True:
            attempts += 1

            try:
                return super().__iter__()
            except OperationalError as ex:
                if 'server closed the connection unexpectedly' not in str(ex):
                    raise

                if attempts <= self.__max_retry_count__:
                    sleep_for = 2 ** (attempts - 1)
                    logging.error(
                        ' [!] Database connection error: retrying Strategy => '
                        f'sleeping for {sleep_for}s and will retry '
                        f'(attempt #{attempts} of {self.__max_retry_count__})'
                        f'\nDetailed query impacted: {ex}'
                    )
                    sleep(sleep_for)
                    continue
                else:
                    raise
            except StatementError as ex:
                if (
                    'reconnect until invalid transaction is rolled back'
                    not in str(ex)
                ):
                    raise
                self.session.rollback()


class Query:
    DATABASE_URI_PARAMS: dict = {
        'db_driver',
        'db_host',
        'db_port',
        'db_name',
        'db_user',
        'db_pass',
        'use_charset_utf8'
    }

    def __init__(
        self, engine, use_charset_utf8: bool = False, template_dir: str = None
    ):
        '''
        engine:
        - obj: SQLAlchemy's engine instance.
        - str: Build SQLAlchemy's engine from string database uri.
        - dict: Build SQLAlchemy's engine from database_uri's func kwargs and
        SSHTunnelForwarder's func kwarg.

        use_charset_utf8: bool
        - Use use_charset_utf8=True for mysql driver if needed.

        template_dir: str
        - Use for set specify jinja2 template directory, os.getcwd()
        value is used as default.
        '''

        self.__in_context = False
        self.__template_dir = template_dir if template_dir else getcwd()
        self.__env = Environment(
            loader=FileSystemLoader(self.__template_dir)
        )
        self.__use_charset_utf8 = use_charset_utf8
        self.__engine = engine
        self.__ssh_param = (
            {
                key: value
                for key, value in self.__engine.items()
                if key not in self.DATABASE_URI_PARAMS
            }
            if isinstance(self.__engine, dict)
            else {}
        )
        self.__db_host = None
        self.__db_port = None

    @property
    def __create_engine(self) -> Engine:
        engine = copy(self.__engine)

        if isinstance(engine, dict):
            engine = {
                key: value
                for key, value in engine.items()
                if key in self.DATABASE_URI_PARAMS
            }

            if self.__use_charset_utf8:
                engine['use_charset_utf8'] = self.__use_charset_utf8

            if all([
                self.__db_host,
                self.__db_port
            ]):
                engine.update({
                    'db_host': self.__db_host,
                    'db_port': self.__db_port
                })

            # Process to returning string database uri.
            url = database_uri(**engine)
        elif isinstance(engine, Engine):
            url = str(engine.url)
        else:
            url = engine

        # "set character set" in sqlalchemy?
        # https://groups.google.com/forum/#!topic/sqlalchemy/3kiPusCy8FM
        if (
            url.startswith('mysql')
            and 'charset=utf8' not in url
            and self.__use_charset_utf8
        ):
            # Add params to given URL in Python
            # https://stackoverflow.com/questions/2506379/add-params-to-given-url-in-python
            (
                query := dict(
                    parse_qsl(
                        (
                            url := list(
                                urlparse(url)
                            )
                        )[4]
                    )
                )
            ).update({'charset': 'utf8'})
            url[4] = urlencode(query)
            url = urlunparse(url)

        return create_engine(
            url,
            json_serializer=json_serializer,
            pool_size=10,
            max_overflow=2,
            pool_recycle=300,
            pool_pre_ping=True,
            pool_use_lifo=True
        )

    def __create_session(self, engine: Engine) -> ScopedSessionMixin:
        # scoped_session(sessionmaker()) or plain sessionmaker() in sqlalchemy?
        # https://stackoverflow.com/questions/6519546/scoped-sessionsessionmaker-or-plain-sessionmaker-in-sqlalchemy
        # how to fix “OperationalError: (psycopg2.OperationalError) server
        # closed the connection unexpectedly”
        # https://stackoverflow.com/questions/55457069/how-to-fix-operationalerror-psycopg2-operationalerror-server-closed-the-conn
        return scoped_session(
            sessionmaker(bind=engine, query_cls=RetryingQuery)
        )()

    def __call__(self, string, **kwargs) -> Iterable[dict]:
        '''
        string: str
        - Use for put raw query sql or sql file path. It is support jinja2
        templating.

        kwargs: dict
        - Put sql parameter and or jinja2 variable in here.
        - Optional parameter can be use for best query result:
            1. generator_mode: bool
                - Use generator_mode=True if we want result as generator.
            2. json_mode: bool
                - Use json_mode=True if we want result as string dumps json.
            3. debug_mode: bool
                - Use debug_mode=False if you don't want to see error
                information.
            4. fetch_size: int
                - Use for set how many rows use on every fetch.
                - Default value is 100,000.
            5. stream_results: bool
                - Use stream_results=False for update/insert/delete query.
        '''

        # Issue with a python function returning a generator or a normal
        # object
        # https://stackoverflow.com/questions/25313283/issue-with-a-python-function-returning-a-generator-or-a-normal-object
        generator_mode = kwargs.pop('generator_mode', None)

        result = self.__query(string, **kwargs)

        if not generator_mode:
            result = (
                ''.join(result)
                if kwargs.get('json_mode') is True
                else list(result)
            )

        return result

    def __compile(self, text_clause):
        '''
        Function to compile and bind parameter sql.

        text_clause: sqlalchemy.sql.elements.TextClause
        '''

        # SQLAlchemy: print the actual query
        # https://stackoverflow.com/questions/5631078/sqlalchemy-print-the-actual-query
        class StringLiteral(String):
            '''
            Teach SA how to literalize various things.
            '''

            def literal_processor(self, dialect):
                super_processor = super(
                    StringLiteral, self
                ).literal_processor(dialect)

                def process(value):
                    if (
                        isinstance(
                            value := serialize(value),
                            Iterable
                        )
                        and not isinstance(value, str)
                    ):
                        value = dumps(value)

                    if isinstance(
                        result := super_processor(value),
                        bytes
                    ):
                        result = result.decode(dialect.encoding)

                    return result

                return process

        class LiteralDialect(DefaultDialect):
            colspecs = {
                DateTime: StringLiteral,
                Date: StringLiteral,
                JSON: StringLiteral
            }

        return str(
            text_clause.compile(
                dialect=LiteralDialect(),
                compile_kwargs={'literal_binds': True}
            )
        )

    def render(self, string: str, literal_binds=True, **kwargs):
        '''
        Function for render jinja2 template and bind parameter sql.

        string: str
        - Use for put raw query sql or sql file path. It is support jinja2
        templating.

        literal_binds: bool
        - Use literal_binds=False if not want to see sql bind parameter.

        kwargs: dict
        - Put sql parameter and or jinja2 variable in here.
        - It is mandatory if string have sql parameter and literal_binds=True.
        '''

        # SQLAlchemy: print the actual query
        # https://stackoverflow.com/questions/5631078/sqlalchemy-print-the-actual-query

        def __template_source(template_file):
            '''
            Function for returning jinja2 template source.

            template_file: str
            - Template file path.
            '''

            return self.__env.loader.get_source(self.__env, template_file)[0]

        def jinja2_keys(template_source):
            # Jinja2 load templates from separate location than working
            # directory
            # https://stackoverflow.com/questions/37968787/jinja2-load-templates-from-separate-location-than-working-directory
            # Template Designer Documentation
            # https://jinja.palletsprojects.com/en/2.11.x/templates/
            # How to get list of all variables in jinja 2 templates
            # https://stackoverflow.com/questions/8260490/how-to-get-list-of-all-variables-in-jinja-2-templates
            # The Meta API
            # https://jinja.palletsprojects.com/en/2.11.x/api/
            (
                keys := find_undeclared_variables(
                    parsed_content := self.__env.parse(template_source)
                )
            ).update({
                key
                for ref_template in find_referenced_templates(
                    parsed_content
                )
                for key in jinja2_keys(
                    __template_source(ref_template)
                )
            })
            return keys

        kwargs = {
            key: value
            for key, value in kwargs.items()
            if key not in [
                'json_mode',
                'debug_mode',
                'fetch_size',
                'stream_results'
            ]
        }

        if isinstance(string, str):
            if isfile(
                path_join(self.__template_dir, string)
            ):
                template = self.__env.get_template(string)
                template_source = __template_source(string)
            else:
                template = self.__env.from_string(string)
                template_source = string

            string = template.render(**kwargs)

            for key in (
                jinja2_keys(template_source).intersection(
                    kwargs.keys()
                ) - set(
                    re.findall(
                        r':(\w+)',
                        string
                    )
                )
            ):
                kwargs.pop(key, None)

        args = []

        for key in kwargs:
            param_kwargs = {
                'key': key,
                'type_': String
            }

            if (
                value := kwargs[key]
            ) is None:
                # How to insert NULL value in SQLAlchemy?
                # https://stackoverflow.com/questions/32959336/how-to-insert-null-value-in-sqlalchemy
                kwargs[key] = null()

            for parameter_type, database_column_type in [
                [type(None), NullType],
                [float, Float],
                [int, Integer],
                [bool, Boolean],
                [date, Date],
                [datetime, DateTime],
                [dict, JSON],
                [Iterable, JSON]
            ]:
                if (
                    isinstance(value, parameter_type)
                    and not isinstance(value, str)
                ):
                    param_kwargs['type_'] = database_column_type
                    break

            args.append(
                bindparam(**param_kwargs)
            )

        string = text(string).bindparams(*args, **kwargs)

        if literal_binds:
            string = self.__compile(string)
        elif literal_binds is not None:
            string = str(string)

        return string

    def __enter__(self):
        self.__in_context = True

        if self.__ssh_param:
            # Python SSHTunnel w/ Paramiko - CLI works, but
            # not in script
            # https://stackoverflow.com/questions/39945269/python-sshtunnel-w-paramiko-cli-works-but-not-in-script
            # Menjalankan query dengan mode SSH
            # SSHTunnelForwarder's func kwargs
            # https://sshtunnel.readthedocs.io/en/latest/#api
            # Setup a SSH Tunnel With the Sshtunnel Module in Python
            # https://blog.ruanbekker.com/blog/2018/04/23/setup-a-ssh-tunnel-with-the-sshtunnel-module-in-python/
            # SSHTunnelForwarder.daemon_forward_servers is not
            # respected:
            # https://github.com/pahaz/sshtunnel/issues/102
            # tunnel without clause:
            # tunnel.daemon_forward_servers = True
            # tunnel.daemon_transport = True
            # tunnel.start()
            # tunnel.stop()
            self.__tunnel = SSHTunnelForwarder(**self.__ssh_param)
            self.__tunnel.__enter__()
            self.__db_host = self.__tunnel.local_bind_host
            self.__db_port = self.__tunnel.local_bind_port

        self.__engine_for_process = self.__create_engine
        self.__session = self.__create_session(self.__engine_for_process)
        return self

    def __exit__(self, *exc):
        self.__in_context = False
        self.__session.commit()
        self.__session.close()
        self.__engine_for_process.dispose()

        if self.__ssh_param:
            self.__tunnel.__exit__(self, *exc)
            self.__db_host = None
            self.__db_port = None

    def __stream_result(self, string: str) -> bool:
        for s in re.findall(
            r"'[^']*'",
            _string := '\n'.join([
                s
                for s in re.sub(
                    r'--.*', '', str(string)
                ).split('\n')
                if s
            ]),
            flags=re.DOTALL
        ):
            _string = re.sub(
                s.replace('[', r'\[').replace(']', r'\]').replace(
                    '(', r'\('
                ).replace(')', r'\)'),
                "''",
                _string,
                1
            )

        regex_sql_space = r'\s+(.*\s+)?'

        # Chapter 13 SQL Statements
        # https://dev.mysql.com/doc/refman/5.6/en/sql-statements.html
        for regex in [
            'EXPLAIN',
            'INSERT',
            'SET',
            regex_sql_space.join(['DELETE', 'FROM']),
            regex_sql_space.join(['MERGE', 'INTO', 'USING']),
            'CREATE',
            'ALTER',
            'DROP',
            regex_sql_space.join(['RENAME', 'TABLE']),
            regex_sql_space.join(['TRUNCATE', 'TABLE']),
            'SHOW',
            'DESCRIBE',
            'CALL',
            'DO',
            'HANDLER',
            'LOAD',
            'REPLACE',
            'KILL',
            'GRANT'
        ]:
            if re.findall(
                fr'\s*{ regex }\s+.*',
                _string,
                flags=re.DOTALL + re.IGNORECASE
            ):
                stream_results = False
                break
        else:
            stream_results = True

        return stream_results

    def __query(self, string, **kwargs) -> Generator:
        '''
        Func to get sql query result.

        string: str
        - Use for put raw query sql or sql file path. It is support jinja2
        templating.

        kwargs: dict
        - Put sql parameter and or jinja2 variable in here.
        - Optional parameter can be use for best query result:
            1. json_mode: bool
                - Use json_mode=True if we want result as string dumps json.
            2. debug_mode: bool
                - Use debug_mode=False if you don't want to see error
                information.
            3. fetch_size: int
                - Use for set how many rows use on every fetch.
                - Default value is 100,000.
            4. stream_results: bool
                - Use stream_results=False for update/insert/delete query.
        '''

        def method(string, **kwargs) -> Generator:
            string = self.render(string, None, **kwargs)

            if not (stream_results := kwargs.get('stream_results', False)):
                stream_results = self.__stream_result(string)

            try:
                query_result = self.__session.execute(
                    string,
                    execution_options={'stream_results': stream_results}
                )

                # memory-efficient built-in SqlAlchemy iterator /
                # generator:
                # https://stackoverflow.com/questions/7389759/memory-efficient-built-in-sqlalchemy-iterator-generator
                # Python: Using Flask to stream chunked dynamic content to end
                # users
                # https://fabianlee.org/2019/11/18/python-using-flask-to-stream-chunked-dynamic-content-to-end-users/
                # Streaming Contents
                # https://flask.palletsprojects.com/en/1.1.x/patterns/streaming/#basic-usage
                # Streaming JSON with Flask
                # https://blog.al4.co.nz/2016/01/streaming-json-with-flask/
                if (json_mode := kwargs.get('json_mode') is True):
                    yield '['

                while True:
                    batch = query_result.fetchmany(
                        kwargs.get('fetch_size', 100_000)
                    ) if query_result.returns_rows else None

                    def to_result(row, json_mode=False):
                        data = dict(row)
                        return (
                            json_serializer(data)
                            if json_mode else deserialize(data)
                        )

                    if not batch:
                        if not query_result.returns_rows:
                            yield to_result(
                                {
                                    'affected_row': query_result.rowcount,
                                    'query': self.__compile(string),
                                    'finish_time': datetime.now()
                                },
                                json_mode
                            )

                        break

                    rows = batch.__iter__()

                    try:
                        prev_row = next(rows)

                        for row in rows:
                            result = to_result(prev_row, json_mode)
                            prev_row = row
                            yield result + ', ' if json_mode else result

                        yield to_result(prev_row, json_mode)
                    except StopIteration:
                        pass
                    # KeyboardInterrupt and SystemExit should not be wrapped by
                    # sqlalchemy #689
                    # https://github.com/sqlalchemy/sqlalchemy/issues/689
                    # Avoiding accidentally catching KeyboardInterrupt and
                    # SystemExit in Python 2.4
                    # https://stackoverflow.com/questions/2669750/avoiding-accidentally-catching-keyboardinterrupt-and-systemexit-in-python-2-4
                    # Except block handles 'BaseException'
                    # https://lgtm.com/rules/6780080/
                    # Catch multiple exceptions in one line (except block)
                    # https://stackoverflow.com/questions/6470428/catch-multiple-exceptions-in-one-line-except-block
                    # How to get the process ID to kill a nohup process?
                    # https://stackoverflow.com/questions/17385794/how-to-get-the-process-id-to-kill-a-nohup-process
                    # Fixing “Lock wait timeout exceeded; try restarting
                    # transaction” for a 'stuck" Mysql table?
                    # https://stackoverflow.com/questions/2766785/fixing-lock-wait-timeout-exceeded-try-restarting-transaction-for-a-stuck-my/10315184
                    except (KeyboardInterrupt, SystemExit, Exception):
                        break

                if json_mode:
                    yield ']'

                query_result.close()
            except (KeyboardInterrupt, SystemExit, Exception):
                self.__session.rollback()

                if kwargs.get('debug_mode') in [None, True]:
                    raise

        if self.__in_context:
            yield from method(string, **kwargs)
        else:
            with self:
                yield from method(string, **kwargs)

    def insert(
        self,
        data: Iterable[dict],
        table_name: str,
        generator_mode=False
    ) -> Iterable[dict]:
        '''
        Bulk Insert (PostgreSQL Only!)

        generator_mode: bool
        - Use generator_mode=True if we want result as generator.
        '''

        def method(data: Iterable[dict], table_name: str) -> Generator:
            count = 0
            csv_input = StringIO()

            for i, row in enumerate(data):
                if i == 0:
                    csv_writer = csv.DictWriter(csv_input, row.keys())
                    csv_writer.writeheader()

                csv_writer.writerow(row)
                count += 1
                yield row
            else:
                csv_input.seek(0)

            if count:
                with self.__session.connection().connection.cursor() as cur:
                    if hasattr(cur, 'copy_expert'):
                        # PostgreSQL
                        cur.copy_expert(
                            f'''
                            COPY
                                {table_name}
                            FROM
                                STDIN
                            WITH (
                                FORMAT CSV,
                                HEADER TRUE
                            )
                            ;
                            ''',
                            csv_input
                        )
                    else:
                        logging.error(message := 'Working only on PostgreSQL!')
                        raise Exception(message)

            logging.debug(
                f'{count:,} row{"s" if count > 1 else ""} '
                f'has been inserted to {table_name}'
            )

        kwargs = {'data': data, 'table_name': table_name}

        if self.__in_context:
            result = method(**kwargs)
        else:
            with self:
                result = method(**kwargs)

        return result if generator_mode else list(result)
