"""drivers/relational_meta_store.py

Base class for relational database drivers, using SQLAlchemy.
"""

import contextlib
import functools
import json
import re
import urllib.parse as urlparse
from abc import ABC, abstractmethod
from collections import OrderedDict
from typing import Any, Dict, Iterator, Mapping, Optional, Sequence, Tuple, Type, Union

import numpy as np
import sqlalchemy as sqla
import terracotta
from sqlalchemy.engine.base import Connection
from sqlalchemy.engine.url import URL
from terracotta import exceptions
from terracotta.drivers.base_classes import (
    KeysType,
    MetaStore,
    MultiValueKeysType,
    requires_writable,
)
from terracotta.profile import trace

_ERROR_ON_CONNECT = (
    "Could not connect to database. Make sure that the given path points "
    "to a valid Terracotta database, and that you ran driver.create()."
)

DATABASE_DRIVER_EXCEPTIONS_TO_CONVERT: Tuple[Type[Exception], ...] = (
    sqla.exc.OperationalError,
    sqla.exc.InternalError,
    sqla.exc.ProgrammingError,
    sqla.exc.InvalidRequestError,
)

ExceptionType = Union[Type[Exception], Tuple[Type[Exception], ...]]


@contextlib.contextmanager
def convert_exceptions(
    error_message: str,
    exceptions_to_convert: ExceptionType = DATABASE_DRIVER_EXCEPTIONS_TO_CONVERT,
) -> Iterator:
    try:
        yield
    except exceptions_to_convert as exception:
        raise exceptions.InvalidDatabaseError(error_message) from exception


@contextlib.contextmanager
def try_database_operation(connection: Connection) -> Iterator:
    try:
        yield
    except:  # noqa: E722
        connection.rollback()
        raise
    else:
        connection.commit()


class RelationalMetaStore(MetaStore, ABC):
    SQL_DIALECT: str  # The database flavour, eg mysql, sqlite, etc
    SQL_DRIVER: str  # The actual database driver, eg pymysql, sqlite3, etc
    SQL_KEY_SIZE: int
    SQL_TIMEOUT_KEY: str

    SQLA_STRING: Any = sqla.types.String
    SQLA_METADATA_TYPE_LOOKUP: Dict[str, Any] = {
        "real": functools.partial(sqla.types.Float, precision=8),
        "text": sqla.types.Text,
        "blob": sqla.types.LargeBinary,
    }

    _METADATA_COLUMNS: Tuple[Tuple[str, str], ...] = (
        ("bounds_north", "real"),
        ("bounds_east", "real"),
        ("bounds_south", "real"),
        ("bounds_west", "real"),
        ("convex_hull", "text"),
        ("valid_percentage", "real"),
        ("min", "real"),
        ("max", "real"),
        ("mean", "real"),
        ("stdev", "real"),
        ("percentiles", "blob"),
        ("metadata", "text"),
    )

    def __init__(self, path: str) -> None:
        settings = terracotta.get_settings()
        db_connection_timeout: int = settings.DB_CONNECTION_TIMEOUT

        self.url = self._parse_path(path)
        self.sqla_engine = sqla.create_engine(
            self.url,
            echo=False,
            future=True,
            connect_args={self.SQL_TIMEOUT_KEY: db_connection_timeout},
            # automatically re-spawn stale connections, see terracotta#266
            pool_pre_ping=True,
        )
        self.sqla_metadata = sqla.MetaData()

        self._db_keys: Optional[OrderedDict] = None

        self._connection: Optional[Connection] = None
        self.connected: bool = False
        self.db_version_verified: bool = False

        # use normalized path to make sure username and password don't leak into __repr__
        super().__init__(self._normalize_path(path))

    @classmethod
    def _parse_path(cls, connection_string: str) -> URL:
        if "//" not in connection_string:
            connection_string = f"//{connection_string}"

        con_params = urlparse.urlparse(connection_string)

        if not con_params.scheme:
            con_params = urlparse.urlparse(f"{cls.SQL_DIALECT}:{connection_string}")

        if con_params.scheme != cls.SQL_DIALECT:
            raise ValueError(f'unsupported URL scheme "{con_params.scheme}"')

        settings = terracotta.get_settings()
        url = URL.create(
            drivername=f"{cls.SQL_DIALECT}+{cls.SQL_DRIVER}",
            username=con_params.username or settings.SQL_USER,
            password=con_params.password or settings.SQL_PASSWORD,
            host=con_params.hostname,
            port=con_params.port,
            database=con_params.path[1:],  # remove leading '/' from urlparse
            query=dict(urlparse.parse_qsl(con_params.query)),
        )

        return url

    @contextlib.contextmanager
    def connect(self, verify: bool = True) -> Iterator:
        @convert_exceptions(_ERROR_ON_CONNECT, sqla.exc.OperationalError)
        def get_connection() -> Connection:
            return self.sqla_engine.connect().execution_options(
                isolation_level="READ UNCOMMITTED"
            )

        if not self.connected:
            # start a new connection
            try:
                with get_connection() as connection:
                    self._connection = connection
                    self.connected = True
                    if verify:
                        self._connection_callback()

                    with try_database_operation(self._connection):
                        yield self._connection
            finally:
                self._connection = None
                self.connected = False
        else:
            # re-use existing connection
            assert self._connection is not None
            with try_database_operation(self._connection):
                yield self._connection

    def _connection_callback(self) -> None:
        if not self.db_version_verified:
            # check for version compatibility
            def version_tuple(version_string: str) -> Sequence[str]:
                return version_string.split(".")

            db_version = self.db_version
            current_version = terracotta.__version__

            if version_tuple(db_version)[:2] != version_tuple(current_version)[:2]:
                raise exceptions.InvalidDatabaseError(
                    f"Version conflict: database was created in v{db_version}, "
                    f"but this is v{current_version}"
                )
            self.db_version_verified = True

    @property
    @convert_exceptions(_ERROR_ON_CONNECT)
    def db_version(self) -> str:
        """Terracotta version used to create the database"""
        terracotta_table = sqla.Table(
            "terracotta", self.sqla_metadata, autoload_with=self.sqla_engine
        )
        stmt = sqla.select(terracotta_table.c.version)

        with self.connect() as conn:
            version = conn.execute(stmt).scalar()

        return str(version)

    @requires_writable
    @convert_exceptions("Could not create database")
    def create(
        self, keys: Sequence[str], key_descriptions: Optional[Mapping[str, str]] = None
    ) -> None:
        """Create and initialize database with empty tables.

        This must be called before opening the first connection. The database must not
        exist already.

        Arguments:

            keys: Key names to use throughout the Terracotta database.
            key_descriptions: Optional (but recommended) full-text description for some keys,
                in the form of ``{key_name: description}``.

        """
        self._create_database()
        self._initialize_database(keys, key_descriptions)

    @abstractmethod
    def _create_database(self) -> None:
        # Note that some subclasses may not actually create any database here, as
        # it may be created automatically on connection for some database vendors
        pass

    def _initialize_database(
        self, keys: Sequence[str], key_descriptions: Optional[Mapping[str, str]] = None
    ) -> None:
        if key_descriptions is None:
            key_descriptions = {}
        else:
            key_descriptions = dict(key_descriptions)

        if not all(k in keys for k in key_descriptions.keys()):
            raise exceptions.InvalidKeyError(
                "key description dict contains unknown keys"
            )

        if not all(re.match(r"^\w+$", key) for key in keys):
            raise exceptions.InvalidKeyError("key names must be alphanumeric")

        if any(key in self._RESERVED_KEYS for key in keys):
            raise exceptions.InvalidKeyError(
                f"key names cannot be one of {self._RESERVED_KEYS!s}"
            )

        terracotta_table = sqla.Table(
            "terracotta",
            self.sqla_metadata,
            sqla.Column("version", self.SQLA_STRING(255), primary_key=True),
        )
        key_names_table = sqla.Table(
            "key_names",
            self.sqla_metadata,
            sqla.Column(
                "key_name", self.SQLA_STRING(self.SQL_KEY_SIZE), primary_key=True
            ),
            sqla.Column("description", self.SQLA_STRING(8000)),
            sqla.Column("idx", sqla.types.Integer, unique=True),
        )
        _ = sqla.Table(
            "datasets",
            self.sqla_metadata,
            *[
                sqla.Column(key, self.SQLA_STRING(self.SQL_KEY_SIZE), primary_key=True)
                for key in keys
            ],
            sqla.Column("path", self.SQLA_STRING(8000)),
        )
        _ = sqla.Table(
            "metadata",
            self.sqla_metadata,
            *[
                sqla.Column(key, self.SQLA_STRING(self.SQL_KEY_SIZE), primary_key=True)
                for key in keys
            ],
            *[
                sqla.Column(name, self.SQLA_METADATA_TYPE_LOOKUP[column_type]())
                for name, column_type in self._METADATA_COLUMNS
            ],
        )
        self.sqla_metadata.create_all(self.sqla_engine)

        with self.connect(verify=False) as conn:
            conn.execute(
                terracotta_table.insert().values(version=terracotta.__version__)
            )
            conn.execute(
                key_names_table.insert(),
                [
                    dict(key_name=key, description=key_descriptions.get(key, ""), idx=i)
                    for i, key in enumerate(keys)
                ],
            )

    @convert_exceptions("Could not retrieve keys from database")
    def get_keys(self) -> OrderedDict:
        keys_table = sqla.Table(
            "key_names", self.sqla_metadata, autoload_with=self.sqla_engine
        )

        with self.connect() as conn:
            result = conn.execute(
                sqla.select(
                    keys_table.c["key_name"], keys_table.c["description"]
                ).order_by(keys_table.c["idx"])
            )
        return OrderedDict((row.key_name, row.description) for row in result.all())

    @property
    def key_names(self) -> Tuple[str, ...]:
        """Names of all keys defined by the database"""
        if self._db_keys is None:
            self._db_keys = self.get_keys()
        return tuple(self._db_keys.keys())

    @trace("get_datasets")
    @convert_exceptions("Could not retrieve datasets")
    def get_datasets(
        self,
        where: Optional[MultiValueKeysType] = None,
        page: int = 0,
        limit: Optional[int] = None,
    ) -> Dict[Tuple[str, ...], str]:
        if where is None:
            where = {}

        where = {
            key: value if isinstance(value, list) else [value]
            for key, value in where.items()
        }

        datasets_table = sqla.Table(
            "datasets", self.sqla_metadata, autoload_with=self.sqla_engine
        )
        stmt = (
            datasets_table.select()
            .where(
                *[
                    sqla.or_(*[datasets_table.c[column] == value for value in values])
                    for column, values in where.items()
                ]
            )
            .order_by(*datasets_table.c.values())
        )

        if limit is not None:
            stmt = stmt.limit(limit).offset(page * limit if limit is not None else None)

        with self.connect() as conn:
            result = conn.execute(stmt).all()

        def keytuple(row: sqla.engine.row.Row) -> Tuple[str, ...]:
            return tuple(getattr(row, key) for key in self.key_names)

        datasets = {keytuple(row): row.path for row in result}
        return datasets

    @trace("get_metadata")
    @convert_exceptions("Could not retrieve metadata")
    def get_metadata(self, keys: KeysType) -> Optional[Dict[str, Any]]:
        metadata_table = sqla.Table(
            "metadata", self.sqla_metadata, autoload_with=self.sqla_engine
        )
        stmt = metadata_table.select().where(
            *[metadata_table.c[key] == value for key, value in keys.items()]
        )

        with self.connect() as conn:
            row = conn.execute(stmt).first()
        if not row:
            return None

        data_columns, _ = zip(*self._METADATA_COLUMNS)
        encoded_data = {col: getattr(row, col) for col in self.key_names + data_columns}
        return self._decode_data(encoded_data)

    @trace("insert")
    @requires_writable
    @convert_exceptions("Could not write to database")
    def insert(
        self, keys: KeysType, path: str, *, metadata: Optional[Mapping[str, Any]] = None
    ) -> None:
        datasets_table = sqla.Table(
            "datasets", self.sqla_metadata, autoload_with=self.sqla_engine
        )
        metadata_table = sqla.Table(
            "metadata", self.sqla_metadata, autoload_with=self.sqla_engine
        )

        with self.connect() as conn:
            conn.execute(
                datasets_table.delete().where(
                    *[
                        datasets_table.c[column] == value
                        for column, value in keys.items()
                    ]
                )
            )
            conn.execute(datasets_table.insert().values(**keys, path=path))

            if metadata is not None:
                encoded_data = self._encode_data(metadata)
                conn.execute(
                    metadata_table.delete().where(
                        *[
                            metadata_table.c[column] == value
                            for column, value in keys.items()
                        ]
                    )
                )
                conn.execute(metadata_table.insert().values(**keys, **encoded_data))

    @trace("delete")
    @requires_writable
    @convert_exceptions("Could not write to database")
    def delete(self, keys: KeysType) -> None:
        if not self.get_datasets(keys):
            raise exceptions.DatasetNotFoundError(f"No dataset found with keys {keys}")

        datasets_table = sqla.Table(
            "datasets", self.sqla_metadata, autoload_with=self.sqla_engine
        )
        metadata_table = sqla.Table(
            "metadata", self.sqla_metadata, autoload_with=self.sqla_engine
        )

        with self.connect() as conn:
            conn.execute(
                datasets_table.delete().where(
                    *[
                        datasets_table.c[column] == value
                        for column, value in keys.items()
                    ]
                )
            )
            conn.execute(
                metadata_table.delete().where(
                    *[
                        metadata_table.c[column] == value
                        for column, value in keys.items()
                    ]
                )
            )

    @staticmethod
    def _encode_data(decoded: Mapping[str, Any]) -> Dict[str, Any]:
        """Transform from internal format to database representation"""
        encoded = {
            "bounds_north": decoded["bounds"][0],
            "bounds_east": decoded["bounds"][1],
            "bounds_south": decoded["bounds"][2],
            "bounds_west": decoded["bounds"][3],
            "convex_hull": json.dumps(decoded["convex_hull"]),
            "valid_percentage": decoded["valid_percentage"],
            "min": decoded["range"][0],
            "max": decoded["range"][1],
            "mean": decoded["mean"],
            "stdev": decoded["stdev"],
            "percentiles": np.array(decoded["percentiles"], dtype="float32").tobytes(),
            "metadata": json.dumps(decoded["metadata"]),
        }
        return encoded

    @staticmethod
    def _decode_data(encoded: Mapping[str, Any]) -> Dict[str, Any]:
        """Transform from database format to internal representation"""
        decoded = {
            "bounds": tuple(
                [encoded[f"bounds_{d}"] for d in ("north", "east", "south", "west")]
            ),
            "convex_hull": json.loads(encoded["convex_hull"]),
            "valid_percentage": encoded["valid_percentage"],
            "range": (encoded["min"], encoded["max"]),
            "mean": encoded["mean"],
            "stdev": encoded["stdev"],
            "percentiles": np.frombuffer(
                encoded["percentiles"], dtype="float32"
            ).tolist(),
            "metadata": json.loads(encoded["metadata"]),
        }
        return decoded
