# Copyright (c) 2022 Tulir Asokan
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
from __future__ import annotations

from typing import Any, AsyncContextManager, Type
from abc import ABC, abstractmethod
from contextlib import asynccontextmanager
from contextvars import ContextVar
import logging

from yarl import URL

from mautrix import __optional_imports__
from mautrix.util.logging import TraceLogger

from .connection import LoggingConnection
from .errors import DatabaseNotOwned, ForeignTablesFound
from .scheme import Scheme
from .upgrade import UpgradeTable, upgrade_tables

if __optional_imports__:
    from aiosqlite import Cursor
    from asyncpg import Record

conn_var: ContextVar[LoggingConnection | None] = ContextVar("db_connection", default=None)


class Database(ABC):
    schemes: dict[str, Type[Database]] = {}
    log: TraceLogger

    scheme: Scheme
    url: URL
    _db_args: dict[str, Any]
    upgrade_table: UpgradeTable | None
    owner_name: str | None
    ignore_foreign_tables: bool

    def __init__(
        self,
        url: URL,
        upgrade_table: UpgradeTable | None,
        db_args: dict[str, Any] | None = None,
        log: TraceLogger | None = None,
        owner_name: str | None = None,
        ignore_foreign_tables: bool = True,
    ) -> None:
        self.url = url
        self._db_args = {**db_args} if db_args else {}
        self.upgrade_table = upgrade_table
        self.owner_name = owner_name
        self.ignore_foreign_tables = ignore_foreign_tables
        self.log = log or logging.getLogger("mau.db")
        assert isinstance(self.log, TraceLogger)

    @classmethod
    def create(
        cls,
        url: str | URL,
        *,
        db_args: dict[str, Any] | None = None,
        upgrade_table: UpgradeTable | str | None = None,
        log: logging.Logger | TraceLogger | None = None,
        owner_name: str | None = None,
        ignore_foreign_tables: bool = True,
    ) -> Database:
        url = URL(url)
        try:
            impl = cls.schemes[url.scheme]
        except KeyError as e:
            if url.scheme in ("postgres", "postgresql"):
                raise RuntimeError(
                    f"Unknown database scheme {url.scheme}."
                    " Perhaps you forgot to install asyncpg?"
                ) from e
            elif url.scheme in ("sqlite", "sqlite3"):
                raise RuntimeError(
                    f"Unknown database scheme {url.scheme}."
                    " Perhaps you forgot to install aiosqlite?"
                ) from e
            raise RuntimeError(f"Unknown database scheme {url.scheme}") from e
        if isinstance(upgrade_table, str):
            upgrade_table = upgrade_tables[upgrade_table]
        elif upgrade_table is None:
            upgrade_table = UpgradeTable()
        elif not isinstance(upgrade_table, UpgradeTable):
            raise ValueError(f"Can't use {type(upgrade_table)} as the upgrade table")
        return impl(
            url,
            db_args=db_args,
            upgrade_table=upgrade_table,
            log=log,
            owner_name=owner_name,
            ignore_foreign_tables=ignore_foreign_tables,
        )

    def override_pool(self, db: Database) -> None:
        pass

    async def start(self) -> None:
        if not self.ignore_foreign_tables:
            await self._check_foreign_tables()
        if self.owner_name:
            await self._check_owner()
        if self.upgrade_table and len(self.upgrade_table.upgrades) > 0:
            await self.upgrade_table.upgrade(self)

    async def _check_foreign_tables(self) -> None:
        if await self.table_exists("state_groups_state"):
            raise ForeignTablesFound("found state_groups_state likely belonging to Synapse")
        elif await self.table_exists("roomserver_rooms"):
            raise ForeignTablesFound("found roomserver_rooms likely belonging to Dendrite")

    async def _check_owner(self) -> None:
        await self.execute(
            """CREATE TABLE IF NOT EXISTS database_owner (
                key   INTEGER PRIMARY KEY DEFAULT 0,
                owner TEXT NOT NULL
            )"""
        )
        owner = await self.fetchval("SELECT owner FROM database_owner WHERE key=0")
        if not owner:
            await self.execute(
                "INSERT INTO database_owner (key, owner) VALUES (0, $1)", self.owner_name
            )
        elif owner != self.owner_name:
            raise DatabaseNotOwned(owner)

    @abstractmethod
    async def stop(self) -> None:
        pass

    @abstractmethod
    def acquire_direct(self) -> AsyncContextManager[LoggingConnection]:
        pass

    @asynccontextmanager
    async def acquire(self) -> LoggingConnection:
        conn = conn_var.get(None)
        if conn is not None:
            yield conn
            return
        async with self.acquire_direct() as conn:
            token = conn_var.set(conn)
            try:
                yield conn
            finally:
                conn_var.reset(token)

    async def execute(self, query: str, *args: Any, timeout: float | None = None) -> str | Cursor:
        async with self.acquire() as conn:
            return await conn.execute(query, *args, timeout=timeout)

    async def executemany(
        self, query: str, *args: Any, timeout: float | None = None
    ) -> str | Cursor:
        async with self.acquire() as conn:
            return await conn.executemany(query, *args, timeout=timeout)

    async def fetch(self, query: str, *args: Any, timeout: float | None = None) -> list[Record]:
        async with self.acquire() as conn:
            return await conn.fetch(query, *args, timeout=timeout)

    async def fetchval(
        self, query: str, *args: Any, column: int = 0, timeout: float | None = None
    ) -> Any:
        async with self.acquire() as conn:
            return await conn.fetchval(query, *args, column=column, timeout=timeout)

    async def fetchrow(
        self, query: str, *args: Any, timeout: float | None = None
    ) -> Record | None:
        async with self.acquire() as conn:
            return await conn.fetchrow(query, *args, timeout=timeout)

    async def table_exists(self, name: str) -> bool:
        async with self.acquire() as conn:
            return await conn.table_exists(name)