import asyncio
import functools
import signal
import time
from collections import defaultdict
from typing import Any, Iterator, Optional

from datadog.dogstatsd.base import statsd
from sqlalchemy.ext.asyncio import AsyncEngine
from structlog import BoundLogger


class ShutdownHandler:
    signal_recieved: bool = False

    def __init__(self, logger: BoundLogger) -> None:
        self.logger = logger
        signal.signal(signal.SIGINT, self.exit_gracefully)
        signal.signal(signal.SIGTERM, self.exit_gracefully)

    def exit_gracefully(self, *args) -> None:
        self.signal_recieved = True
        self.logger.info("Shutdown signal recieved, shutting down")


def deep_chunks(iterable: list[list[Any]], chunk_size: int) -> Iterator[list[Any]]:
    for inner_iterable in iterable:
        yield from chunks(inner_iterable, chunk_size)


def chunks(iterable: list[Any], chunk_size: int) -> Iterator[list[Any]]:
    """Yields slices of iterable object with size of slice limited by chunk_size argument
    >>> chunks([1,2,3,4,5,6,7], 2)
    <generator object chunks at 0x7fbbaa032c00>
    >>> list(chunks(1,2,3,4,5,6,7), 2)
    [[1, 2], [3]]
    """
    for i in range(0, len(iterable), chunk_size):
        yield iterable[i : i + chunk_size]


def ttl_cache(ttl_seconds):
    def decorator(func):
        items = {}
        expiration_times = {}

        @functools.wraps(func)
        def wrapper(*args):
            if args in items:
                if time.time() > expiration_times[args]:
                    # The item has expired, so remove it from the cache.
                    del items[args]
                    del expiration_times[args]
                else:
                    # The item is still valid, so return it from the cache.
                    return items[args]

            # The item is not in the cache, so compute it and add it to the cache.
            result = func(*args)
            items[args] = result
            expiration_times[args] = time.time() + ttl_seconds
            return result

        return wrapper

    return decorator


def async_cache(timeout: Optional[float] = None):
    """Cache imlementation for async functions"""

    def func_decorator(async_func):
        cache = defaultdict(dict)

        async def async_wrapper(*args, **kwargs):
            """Wrapper for async function"""
            key = args
            for kwarg in kwargs.items():
                key += kwarg
            hashed_key = hash(key)

            if cache.get(hashed_key):
                if timeout is not None:
                    if cache[hashed_key]["expired_at"] > time.time():
                        return cache[hashed_key]["value"]
                else:
                    return cache[hashed_key]["value"]

            cache[hashed_key]["value"] = await async_func(*args, **kwargs)
            if timeout is not None:
                cache[hashed_key]["expired_at"] = time.time() + timeout
            return cache[hashed_key]["value"]

        return async_wrapper

    return func_decorator


def async_retry(delay: int = 5, max_retries: int = 5):
    def async_func_handler(func):
        async def async_wrapper(*args, **kwargs):
            retries = 0
            while True:
                try:
                    result = await func(*args, **kwargs)
                    return result
                except Exception:
                    retries += 1
                    if retries > max_retries:
                        raise
                    await asyncio.sleep(delay)

        return async_wrapper

    return async_func_handler


@async_retry(delay=1)
async def async_execute(engine: AsyncEngine, query):
    with statsd.timed(
        "dapd-api-scraper.db_operation_time",
        tags=[
            f"db:{engine.url.database}",
            f"operation:{'read' if query.is_select else 'write'}",
        ],
        use_ms=True,
    ):
        async with engine.connect() as conn:
            return await conn.execute(query)