# SPDX-License-Identifier: GPL-2.0-only
#
# This file is part of Nominatim. (https://nominatim.org)
#
# Copyright (C) 2022 by the Nominatim developer community.
# For a full list of authors see the git log.
"""
Preprocessing of SQL files.
"""
from typing import Set, Dict, Any, cast
import jinja2

from nominatim.db.connection import Connection
from nominatim.db.async_connection import WorkerPool
from nominatim.config import Configuration

def _get_partitions(conn: Connection) -> Set[int]:
    """ Get the set of partitions currently in use.
    """
    with conn.cursor() as cur:
        cur.execute('SELECT DISTINCT partition FROM country_name')
        partitions = set([0])
        for row in cur:
            partitions.add(row[0])

    return partitions


def _get_tables(conn: Connection) -> Set[str]:
    """ Return the set of tables currently in use.
    """
    with conn.cursor() as cur:
        cur.execute("SELECT tablename FROM pg_tables WHERE schemaname = 'public'")

        return set((row[0] for row in list(cur)))

def _get_middle_db_format(conn: Connection, tables: Set[str]) -> str:
    """ Returns the version of the slim middle tables.
    """
    if 'osm2pgsql_properties' not in tables:
        return '1'

    with conn.cursor() as cur:
        cur.execute("SELECT value FROM osm2pgsql_properties WHERE property = 'db_format'")
        row = cur.fetchone()

        return cast(str, row[0]) if row is not None else '1'


def _setup_tablespace_sql(config: Configuration) -> Dict[str, str]:
    """ Returns a dict with tablespace expressions for the different tablespace
        kinds depending on whether a tablespace is configured or not.
    """
    out = {}
    for subset in ('ADDRESS', 'SEARCH', 'AUX'):
        for kind in ('DATA', 'INDEX'):
            tspace = getattr(config, f'TABLESPACE_{subset}_{kind}')
            if tspace:
                tspace = f'TABLESPACE "{tspace}"'
            out[f'{subset.lower()}_{kind.lower()}'] = tspace

    return out


def _setup_postgresql_features(conn: Connection) -> Dict[str, Any]:
    """ Set up a dictionary with various optional Postgresql/Postgis features that
        depend on the database version.
    """
    pg_version = conn.server_version_tuple()
    postgis_version = conn.postgis_version_tuple()
    pg11plus = pg_version >= (11, 0, 0)
    ps3 = postgis_version >= (3, 0)
    return {
        'has_index_non_key_column': pg11plus,
        'spgist_geom' : 'SPGIST' if pg11plus and ps3 else 'GIST'
    }

class SQLPreprocessor:
    """ A environment for preprocessing SQL files from the
        lib-sql directory.

        The preprocessor provides a number of default filters and variables.
        The variables may be overwritten when rendering an SQL file.

        The preprocessing is currently based on the jinja2 templating library
        and follows its syntax.
    """

    def __init__(self, conn: Connection, config: Configuration) -> None:
        self.env = jinja2.Environment(autoescape=False,
                                      loader=jinja2.FileSystemLoader(str(config.lib_dir.sql)))

        db_info: Dict[str, Any] = {}
        db_info['partitions'] = _get_partitions(conn)
        db_info['tables'] = _get_tables(conn)
        db_info['reverse_only'] = 'search_name' not in db_info['tables']
        db_info['tablespace'] = _setup_tablespace_sql(config)
        db_info['middle_db_format'] = _get_middle_db_format(conn, db_info['tables'])

        self.env.globals['config'] = config
        self.env.globals['db'] = db_info
        self.env.globals['postgres'] = _setup_postgresql_features(conn)


    def run_string(self, conn: Connection, template: str, **kwargs: Any) -> None:
        """ Execute the given SQL template string on the connection.
            The keyword arguments may supply additional parameters
            for preprocessing.
        """
        sql = self.env.from_string(template).render(**kwargs)

        with conn.cursor() as cur:
            cur.execute(sql)
        conn.commit()


    def run_sql_file(self, conn: Connection, name: str, **kwargs: Any) -> None:
        """ Execute the given SQL file on the connection. The keyword arguments
            may supply additional parameters for preprocessing.
        """
        sql = self.env.get_template(name).render(**kwargs)

        with conn.cursor() as cur:
            cur.execute(sql)
        conn.commit()


    def run_parallel_sql_file(self, dsn: str, name: str, num_threads: int = 1,
                              **kwargs: Any) -> None:
        """ Execute the given SQL files using parallel asynchronous connections.
            The keyword arguments may supply additional parameters for
            preprocessing.

            After preprocessing the SQL code is cut at lines containing only
            '---'. Each chunk is sent to one of the `num_threads` workers.
        """
        sql = self.env.get_template(name).render(**kwargs)

        parts = sql.split('\n---\n')

        with WorkerPool(dsn, num_threads) as pool:
            for part in parts:
                pool.next_free_worker().perform(part)