import logging
import re
import time
from typing import Any, Optional

import psycopg
from flask import g
from psycopg.abc import Query
from psycopg.connection import Connection
from psycopg.cursor import Cursor
from psycopg.client_cursor import ClientCursor
from psycopg.errors import QueryCanceled
from psycopg.rows import dict_row
from psycopg.types.string import TextLoader
from psycopg_pool import ConnectionPool

from src.config import Configuration
from src.internals.cache.redis import KemonoRedisLock, get_conn
from src.internals.serializers import safe_dumper, safe_loader

pool: Optional[ConnectionPool] = None


def init():
    global pool
    if not pool:
        try:
            pool = ConnectionPool(
                min_size=1,
                max_size=Configuration().webserver["threads"],
                kwargs=dict(
                    host=Configuration().database["host"],
                    dbname=Configuration().database["database"],
                    user=Configuration().database["user"],
                    password=Configuration().database["password"],
                    port=Configuration().database["port"],
                    row_factory=dict_row,
                    application_name=Configuration().database["application_name"],
                    autocommit=True,
                ),
            )
            # citext https://github.com/psycopg/psycopg/issues/227
            with pool.connection() as connection:
                connection.cursor().execute("CREATE EXTENSION IF NOT EXISTS citext;")
                psycopg.types.TypeInfo.fetch(connection, "citext").register()
                psycopg.adapters.register_loader("citext", TextLoader)  # this works for single db case
        except Exception as error:
            logging.exception(f"Failed to connect to the database: {error}")
            raise
    return True


def get_pool() -> ConnectionPool:
    if not pool:
        raise Exception("Database gone")
    return pool


def close_pool():
    global pool
    if pool:
        pool.close()
    pool = None


def get_connection() -> Connection:
    if "connection" not in g:
        g.connection = get_pool().getconn()
    g.connection.autocommit = True
    return g.connection


def get_cursor() -> Cursor:
    if "cursor" not in g:
        if "connection" not in g:
            g.connection = get_pool().getconn()
        g.connection.autocommit = True
        g.cursor = g.connection.cursor()
    return g.cursor


def get_client_binding_cursor() -> ClientCursor:
    if "client_binding_cursor" not in g:
        if "connection" not in g:
            g.connection = get_pool().getconn()
        g.connection.autocommit = True
        g.client_binding_cursor = psycopg.ClientCursor(g.connection)
    return g.client_binding_cursor


def reset_cursor() -> None:
    try:
        g.cursor.execute("ROLLBACK")
    except Exception:  # noqa
        pass
    try:
        g.cursor.close()
    except Exception:  # noqa
        pass
    if "connection" in g:
        if g.connection.closed:
            get_pool().putconn(g.connection)
            g.connection = get_pool().getconn()
            g.cursor = g.connection.cursor()
        else:
            g.cursor = g.connection.cursor()
    else:
        g.connection = get_pool().getconn()
        g.connection.autocommit = True
        g.cursor = g.connection.cursor()


def get_from_cache(redis: Any, cache_key: str, cache_store_method: str):
    if cache_store_method == "set":
        return redis.get(cache_key)
    elif cache_store_method == "rpush":
        result = redis.lrange(cache_key, 0, -1)
        return None if not result else result
    elif cache_store_method == "sadd":
        result = redis.smembers(cache_key)
        return None if not result else result
    raise Exception("Not Implemented")


def cached_query(
    query: str,
    cache_key: str,
    params: tuple | dict = (),
    serialize_fn=safe_dumper,
    deserialize_fn=safe_loader,
    reload: bool = False,
    single: bool = False,
    ex: int = Configuration().redis["default_ttl"],
    ex_on_null: Optional[int] = None,
    prepare: bool = True,
    client_bind: bool = False,
    lock_enabled: bool = Configuration().webserver["use_redis_by_lock_default_on_queries"],
    sets_to_fetch: list[int] | None = None,
    cache_store_method: str = "set",
) -> Any:
    redis = get_conn()
    result = get_from_cache(redis, cache_key, cache_store_method) if not reload else None
    if result is None:
        lock = KemonoRedisLock(redis, cache_key, expire=60, auto_renewal=True)
        was_locked = not lock.acquire(blocking=False) if lock_enabled else False
        if was_locked:
            while not lock.acquire(blocking=False):
                time.sleep(0.1)
            result = get_from_cache(redis, cache_key, cache_store_method) if not reload else None
            if result is not None:
                return deserialize_fn(result)
        try:
            if client_bind:
                cursor = get_client_binding_cursor()
                cursor.execute(query, params, prepare=prepare)
            else:
                cursor = get_cursor()
                cursor.execute(query, params, prepare=prepare)
            if single:
                if sets_to_fetch:
                    if len(sets_to_fetch) == 1:
                        for i in range(sets_to_fetch[0]):
                            cursor.nextset()
                        result = cursor.fetchone()
                    else:
                        result = []
                        for i in range(max(sets_to_fetch)):
                            if i in sets_to_fetch:
                                return result.append(cursor.fetchone())
                            cursor.nextset()
                else:
                    result = cursor.fetchone()  # todo make it "or {}" so we can tell on MGET if either it was in cache
            else:
                if sets_to_fetch:
                    if len(sets_to_fetch) == 1:
                        for i in range(sets_to_fetch[0]):
                            cursor.nextset()
                        result = cursor.fetchall()
                    else:
                        result = []
                        for i in range(max(sets_to_fetch)):
                            if i in sets_to_fetch:
                                return result.append(cursor.fetchall())
                            cursor.nextset()
                else:
                    result = cursor.fetchall()
            if cache_store_method == "set":
                redis.set(
                    cache_key,
                    serialize_fn(result),
                    ex=ex if not ex_on_null or result else ex_on_null,
                )
            elif cache_store_method == "rpush":
                to_push = serialize_fn(result)
                if to_push:
                    redis.rpush(
                        cache_key,
                        *serialize_fn(result),
                    )
                    redis.expire(cache_key, ex)
                else:
                    redis.rpush(
                        cache_key,
                        *serialize_fn([{}]),
                    )
                    redis.expire(cache_key, ex)
                    redis.rpop(cache_key)
            elif cache_store_method == "sadd":
                if result:
                    redis.sadd(
                        cache_key,
                        *serialize_fn(result),
                    )
                    redis.expire(cache_key, ex)
            else:
                getattr(redis, cache_store_method)(
                    cache_key,
                    serialize_fn(result),
                )
        except QueryCanceled:
            reset_cursor()
            raise
        except Exception as query_exception:
            if not re.search(
                r"pgroonga: \[[\w-]+\]\[query\] failed to parse expression:",
                str(query_exception),
            ):
                logging.exception("Failed cached query", extra=dict(query=query, params=params))
            reset_cursor()
            raise
        finally:
            if lock_enabled:
                lock.release()
    else:
        result = deserialize_fn(result)
    return result


def cached_count(
    query: str,
    cache_key: str,
    params: tuple = (),
    reload: bool = False,
    ex: int = Configuration().redis["default_ttl"],
    ex_on_null: Optional[int] = None,
    prepare: bool = True,
    client_bind: bool = False,
    lock_enabled: bool = Configuration().webserver["use_redis_by_lock_default_on_queries"],
    sets_to_fetch: list[int] | None = None,
) -> int:
    # todo: fix serialize_fn
    result = cached_query(
        query,
        cache_key,
        params,
        lambda val: val["count"],
        int,
        reload,
        True,
        ex,
        ex_on_null,
        prepare,
        client_bind,
        lock_enabled,
        sets_to_fetch,
    )
    if isinstance(result, dict):
        result = result["count"]
    return result


def query_db(query: Query, params: tuple | dict = ()) -> list[dict]:
    cursor = get_cursor()
    cursor.execute(query, params)
    result = cursor.fetchall()
    return result


def query_one_db(
    query: Query,
    params: tuple | dict = (),
) -> Optional[dict]:
    cursor = get_cursor()
    cursor.execute(query, params)
    return cursor.fetchone()


def query_rowcount_db(
    query: Query,
    params: tuple = (),
    prepare=True,
) -> int:
    cursor = get_cursor()
    cursor.execute(query, params, prepare=prepare)
    return cursor.rowcount
