"""
File Purpose: misc. tools related to SQL.

Note that many of these solve the same issues that SQLAlchemy solves.
If you already know SQL well, maybe it would be better to use SQLAlchemy.
But, these tools were added so that PlasmaCalcs developers and users
can maybe learn SQL more easily, with less of the SQLAlchemy abstraction.
"""

import os
import contextlib
import logging
import sqlite3
from urllib.parse import urlparse

from .imports import ImportFailed

try:
    import psycopg
except ImportError as err:
    psycopg = ImportFailed("psycopg", err=err, locals=locals(), abbrv='psycopg')
try:
    import sqlglot
except ImportError as err:
    sqlglot = ImportFailed("sqlglot", err=err, locals=locals(), abbrv='sqlglot')

from .math import is_number 
from .sentinels import UNSET
from ..errors import InputError, InputConflictError


### --------------------- Misc helper methods --------------------- ###

def sql_location_string(location=None):
    '''returns "normalized" location string from location input.
    location: None or str
        None or ':memory:' --> in-memory, equivalent to ":memory:". Will not edit any files.
        str --> Online database connection string starts with "postgresql://" or "postgres://".
                All other strings interpreted as filepath, and converted to abspath.
    '''
    if location is None or location == ':memory:':
        return ':memory:'
    elif location.startswith('postgresql://') or location.startswith('postgres://'):
        return location
    else:
        return os.path.abspath(location)

def sql_engine_from_location(location):
    '''returns sql engine implied by location string.
    Possible results:
        "memory": sqlite3 in-memory connection
        "sqlite": sqlite3 connection to file
        "postgres": PostgreSQL connection to online database

    location: str
        implies sql connection, e.g. ":memory:", filepath, or online database connection string.
        postgresql connection string starts with "postgresql://" or "postgres://".

    Currently assumes any non-memory, non-postgres connection is a filepath,
        without checking whether the file exists yet.
    '''
    if location == ':memory:':
        return 'memory'
    elif location.startswith('postgresql://') or location.startswith('postgres://'):
        return 'postgres'
    else:
        return 'sqlite'

def sql_engine_backend(engine):
    '''return sql engine "backend" implied by engine string. Usually equivalent to engine string.
    But, some engines might have same backend, e.g., in-memory and file-based sqlite.

    The implementation here does:
        - 'memory' --> 'sqlite'
        - all other engines unchanged
    '''
    if engine == 'memory':
        return 'sqlite'
    else:
        return engine

def _maybe_assert_sql_location_exists(location, *, assert_exists=None):
    '''depending on assert_exists, may assert location exists or that it does not exist.

    location: str
        implies sql connection, e.g. ":memory:", filepath, or online database connection string.
        if filepath, should be input as abspath.
    assert_exists: None, True, or False
        True --> assert that implied location already exists as a file.
        False --> assert that implied location does not yet exist as a file.
        None --> no assertion is made.
        (Here, ':memory:' never "already exists", since it always means a new database.)
    '''
    if assert_exists is not None:
        if assert_exists is not True and assert_exists is not False:
            raise InputError(f"assert_exists must be None, True, or False, but got {assert_exists!r}")
        engine = sql_engine_from_location(location)
        if engine == 'memory':
            if assert_exists:
                raise InputConflictError("assert_exists=True but location==':memory:'.")
        elif engine == 'sqlite':
            exists = os.path.exists(location)
            if assert_exists is True and not exists:
                raise FileNotFoundError(f"assert_exists=True but database location does not exist: {location!r}")
            elif assert_exists is False and exists:
                raise FileExistsError(f"assert_exists=False but database location already exists: {location!r}")
        elif engine == 'postgres':
            try:
                with psycopg.connect(location):
                    pass
            except psycopg.OperationalError as err:
                if assert_exists is True:
                    errmsg = f"assert_exists=True but failed to connect to database at location: {location!r}."
                    raise FileNotFoundError(errmsg) from err
            else:
                if assert_exists is False:
                    errmsg = f"assert_exists=False but successfully connected to database at location: {location!r}."
                    raise FileExistsError(errmsg)

def sql_to_postgres(sql, *, dialect=UNSET, as_list=False, suppress_warnings=True):
    '''transpile sql statement into postgres.

    dialect: UNSET, None, or str
        The dialect corresponding to the input SQL.
        UNSET --> auto-detect; checks sqlite first, then postgres, then generic sql.
        None --> use generic sql, i.e. sqlglot's default dialect.
        str --> e.g., "sqlite", "postgres". Can be any dialect string recognized by sqlglot.
    as_list: bool
        whether to return result as a list of sql statements.
        False --> return as a single string, with statements separated by semicolons and newlines.
                The last statement will not have a semicolon at the end.
        True --> return a list of sql statements, none of which end with a semicolon.
    suppress_warnings: bool or int
        whether to suppress sqlglot warnings, e.g. about unsupported features of a dialect.
    '''
    try_dialects = ['sqlite', 'postgres', None] if dialect is UNSET else [dialect]
    with _sqlglot_suppress_warnings() if suppress_warnings else contextlib.nullcontext():
        for i, dial in enumerate(try_dialects):
            try:
                result = sqlglot.transpile(sql, read=dial, write='postgres')
            except Exception:
                if i == len(try_dialects) - 1:
                    raise  # last attempt, just raise the error if it fails.
            else:
                break
    return result if as_list else ';\n'.join(result)

@contextlib.contextmanager
def _sqlglot_suppress_warnings():
    '''context manager to suppress sqlglot warnings, e.g. about unsupported features of a dialect.'''
    logger = logging.getLogger("sqlglot")
    old_level = logger.level
    try:
        logger.setLevel(logging.ERROR)
        yield
    finally:
        logger.setLevel(old_level)


### --------------------- Connection Wrappers --------------------- ###

class Sqlite3InMemoryConnection():
    '''class which mimics a database connection to an in-memory database.
    Behaves like sqlite3.connect(':memory:') but will not close() upon __exit__.
    Allows consistent interface with "with" blocks without losing the in-memory connection,
        as long as the InMemoryConnection object itself is kept alive.

    con: None or sqlite3.Connection
        connection to in-memory database.
        If None, creates new sqlite3.Connection to ':memory:'.
    _owner_close_con_on_delete: bool
        whether self should try to close self.con when self is about to be deleted.
        Only applies if self created con, i.e. con not None during init.
    '''
    def __init__(self, *, con=None, _owner_close_con_on_delete=True):
        self._owns_con = con is None
        self._owner_close_con_on_delete = _owner_close_con_on_delete
        if con is None:
            con = sqlite3.connect(':memory:')
        con = sqlite3.connect(':memory:')
        if isinstance(con, Sqlite3InMemoryConnection):
            con = con.con
        self.con = con

    # # # CONTEXT MANAGER # # #
    def __enter__(self):
        '''enter "with" block; returns self (which behaves like sqlite3.Connection)'''
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        '''exit "with" block; does NOT close connection.'''
        pass

    def __del__(self):
        '''delete self; if self owns con and _owner_close_con_on_delete, close con.'''
        if self._owns_con and self._owner_close_con_on_delete:
            try:
                self.con.close()
            except Exception:  # it's fine if con was already closed...
                pass

    # # # OTHER CON-LIKE BEHAVIORS # # #
    def __getattr__(self, name):
        '''pass through any other attribute access to self.con'''
        return getattr(self.con, name)


class PostgresConnection():
    '''class which wraps a database connection to an online PostgreSQL database.
    Allows for sqlite3-like syntax, including executescript(), and using '?' placeholders.

    location: str
        online database connection string, starting with "postgresql://" or "postgres://".
    '''
    def __init__(self, location):
        self.location = location

    # # # CONTEXT MANAGER # # #
    def __enter__(self):
        '''enter "with" block; returns self (which behaves like psycopg connection)'''
        self.con = psycopg.connect(self.location)
        self.con.__enter__()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        '''exit "with" block; closes connection.'''
        self.con.__exit__(exc_type, exc_value, traceback)

    # # # TRANSLATION OF DIALECTS # # #
    sql_to_postgres = staticmethod(sql_to_postgres)

    # # # OTHER CON-LIKE BEHAVIORS # # #
    def __getattr__(self, name):
        '''pass through to self.con any attribute access not defined here.'''
        return getattr(self.con, name)

    def executescript(self, sql):
        '''execute multiple SQL statements at once, separated by semicolons.
        psycopg doesn't have executescript(), so we need to implement it ourselves.
        '''
        statements = sql_to_postgres(sql, as_list=True)
        for s in statements:
            self.con.execute(s)

    def execute(self, query, params=None, **kw_execute):
        '''execute SQL query. Can be written in sqlite-like syntax, e.g. using ? placeholders.'''
        pg_query = self.sql_to_postgres(query)
        return self.con.execute(pg_query, params, **kw_execute)


### --------------------- SqlBase --------------------- ###

class SqlBase():
    '''class which provides some basic SQL-related tools for a database connection.

    location: None or str
        None --> in-memory, equivalent to ":memory:". Will not edit any files.
        str --> connection string. Filepath for local sqlite connection.
                Online database connection string starts with "postgresql://" or "postgres://".
    assert_exists: None, True, or False
        True --> assert that implied location already exists as a file or online database.
        False --> assert that the implied location does not yet exist.
        None --> no assertion is made.
        (Here, ':memory:' never "already exists", since it always means a new database.)
    '''
    memory_connection_cls = Sqlite3InMemoryConnection
    postgres_connection_cls = PostgresConnection

    def __init__(self, location=None, *, assert_exists=None):
        self.location = sql_location_string(location)
        _maybe_assert_sql_location_exists(self.location, assert_exists=assert_exists)

    # # # SQL CONNECTION AND CROSS-ENGINE COMPATIBILITY # # #
    engine = property(lambda self: sql_engine_from_location(self.location),
            doc='''sql engine implied by location string. E.g. 'memory', 'sqlite', or 'postgres'.
            Might be more specific than just an sql dialect. See also: self.engine_backend.''')
    engine_backend = property(lambda self: sql_engine_backend(self.engine),
            doc='''sql engine "backend" implied by engine string. Corresponds to an sql dialect.
            E.g., 'sqlite' or 'postgres', but not 'memory' ('memory' engine <--> 'sqlite' backend).''')

    @property
    def is_local(self):
        '''returns whether self corresponds to a local database connection.
        Doesn't require the connection to actually exist. Currently just checks engine,
            e.g. 'postgres' --> is_local=False; 'memory' or 'sqlite' --> is_local=True.
        '''
        return self.engine in ['memory', 'sqlite']

    def _location_exists(self):
        '''return whether self.location already exists,
        as a cached in-memory connection, a file, or an online database.
        Caution: for online database, always returns True, i.e. assume it exists,
            because this code is not able to create a new online database from scratch.

        Subclass may want to implement self.schema_exists() which uses this method,
            and tells whether the relevant schema has been initialized already.
        '''
        engine = self.engine
        if engine == 'memory':
            return getattr(self, '_in_memory_con', None) is not None
        elif engine == 'sqlite':
            return os.path.exists(self.location)
        elif engine == 'postgres':
            return True

    def connect(self):
        '''return a connection to this database, creating it if it does not exist yet.
        (Cannot create online databases; those will just crash if connection fails.)

        Can be used as a context manager to ensure it is closed properly:
            with sqlbase.connect() as con:
                results = con.execute('SELECT ...').fetchall()
        Else, caller is responsible for closing the connection via con.close().
        '''
        engine = self.engine
        if engine == 'memory':
            # cache in-memory connection, so data persists for the lifetime of this Registry object.
            con = getattr(self, '_in_memory_con', None)
            exists = (con is not None)
            if not exists:
                con = self.memory_connection_cls()
                self._in_memory_con = con
        elif engine == 'sqlite':  # file-based database
            exists = os.path.exists(self.location)
            if not exists:  # makedirs to file if needed.
                os.makedirs(os.path.dirname(self.location), exist_ok=True)
            con = sqlite3.connect(self.location)
        elif engine == 'postgres':  # online database connection
            # assume we can connect, crash if we can't;
            # this code is not empowered to make a new online database.
            con = self.postgres_connection_cls(self.location)
        return con

    def execute(self, *args, **kwargs):
        '''shortcut for with self.connect() as con: return con.execute(*args, **kwargs)'''
        with self.connect() as con:
            return con.execute(*args, **kwargs)

    # # # SANITIZING # # #
    def sanitize_value(self, v):
        '''return sanitized value v, so that it can be stored in a SQL database.
        Sanitized means number, None, or string.
        '''
        if is_number(v):
            return v
        elif v is None:
            return None
        else:
            return str(v)

    def sanitize_values(self, d):
        '''return copy of dict d but with values sanitized to number, None, or string,
        so that they can be stored in a SQL database.
        '''
        sanitized = {}
        for k, v in d.items():
            sanitized[k] = self.sanitize_value(v)
        return sanitized

    # # # INSPECTION # # #
    def get_tables(self, *, _con=None):
        '''return list of all table names in this database.
        _con: None or open connection object
            for efficiency, can provide connection object if already created.
        '''
        with contextlib.ExitStack() as context_stack:
            if _con is None:
                con = self.connect()
                context_stack.enter_context(con)
            else:
                con = _con
            # sqlite3 solution (will need others for other engines):
            if self.engine_backend == 'sqlite':
                tables = con.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall()
                result = [row[0] for row in tables]
            else:
                assert self.engine_backend == 'postgres'
                tables = con.execute('''
                    SELECT table_name FROM information_schema.tables
                    WHERE table_schema = 'public'
                ''').fetchall()
                result = [row[0] for row in tables]
        return result

    @property
    def tables(self):
        '''list of all table names in this database.'''
        return self.get_tables()

    def columns(self, table):
        '''returns list of columns names in this table in this database.'''
        if self.engine_backend == 'sqlite':
            return [row[1] for row in self.execute(f'PRAGMA table_info({table})').fetchall()]
        else:
            assert self.engine_backend == 'postgres'
            columns = self.execute('''
                SELECT column_name FROM information_schema.columns
                WHERE table_name = %s
            ''', (table,)).fetchall()
            return [row[0] for row in columns]

    def column_types(self, table):
        '''returns the types of columns in this table in this database, as str-valued dict.'''
        if self.engine_backend == 'sqlite':
            return {row[1]: row[2] for row in self.execute(f'PRAGMA table_info({table})').fetchall()}
        else:
            assert self.engine_backend == 'postgres'
            columns = self.execute('''
                SELECT column_name, data_type FROM information_schema.columns
                WHERE table_name = %s
            ''', (table,)).fetchall()
            return {row[0]: row[1] for row in columns}

    def text_columns(self, table):
        '''returns the columns in this table whose types are text-based, e.g. TEXT, CHAR(32).'''
        column_types = self.column_types(table)
        text_cols = []
        for col, t in column_types.items():
            t = t.lower()
            for prefix in ['char', 'text', 'varchar']:
                if t.startswith(prefix):
                    text_cols.append(col)
                    break
        return text_cols

    def missing_columns(self, table, columns):
        '''returns list of which of these columns do not yet exist in table in this database.'''
        exists = set(self.columns(table))
        return [col for col in columns if col not in exists]

    def primary_keys(self, table):
        '''returns list of primary key column names in this table in this database.'''
        if self.engine_backend == 'sqlite':
            return [row[1] for row in self.execute(f'PRAGMA table_info({table})').fetchall() if row[5] == 1]
        else:
            assert self.engine_backend == 'postgres'
            pks = self.execute(f'''
                SELECT a.attname
                FROM   pg_index i
                JOIN   pg_attribute a ON a.attrelid = i.indrelid
                                    AND a.attnum = ANY(i.indkey)
                WHERE  i.indrelid = {table!r}::regclass
                AND    i.indisprimary;
            ''').fetchall()
            return [row[0] for row in pks]

    def primary_key(self, table):
        '''returns the single primary key column name in this table in this database,
        crashing if there is not exactly 1 primary key.
        '''
        result = self.primary_keys(table)
        if len(result) != 1:
            raise RegistryContentsError(f"Expected exactly 1 primary key in table {table!r}, but got: {result}")
        return result[0]

    @property
    def sizes(self):
        '''dict of sizes (number of entries) of all tables in this database.'''
        with self.connect() as con:
            tables = self.tables
            result = {}
            for table in tables:
                count = con.execute(f'SELECT COUNT(*) FROM {table}').fetchone()[0]
                result[table] = count
        return result

    # # # DISPLAY # # #
    def _location_display(self):
        '''return self.location but without any sensitive info, suitable for displaying in repr.
        For in-memory or filepath locations, just use self.location directly.
        For online database connection strings, just display engine and database name.
        '''
        if self.engine_backend == 'sqlite':
            return self.location
        else:
            assert self.engine_backend == 'postgres'
            parsed = urlparse(self.location)
            dbname = parsed.path.lstrip('/')
            hostname = parsed.hostname.split('.')[0] if parsed.hostname else 'unknownhost'
            return f'postgresql://{dbname}/host={hostname}'

    def _repr_contents(self):
        '''list of contents for self.__repr__'''
        return [f'location={self._location_display()!r}', f'sizes={self.sizes}']

    def __repr__(self):
        return f'{type(self).__name__}({", ".join(self._repr_contents())})'
