from datetime import datetime, timedelta
from functools import wraps
from io import BytesIO

from fastapi.responses import StreamingResponse

CACHE_SIZE = 50

_cache = {}
_cache_time = {}


def ttl_cache(key_name, media_type=None, ttl_secs=20):
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            # Assuming the prompt is the key for caching, change as necessary
            key = kwargs.get(key_name)
            ttl = timedelta(seconds=ttl_secs)
            # Check cache
            if key in _cache:
                if datetime.now() - _cache_time[key] > ttl:
                    # Cache has expired
                    del _cache[key]
                    del _cache_time[key]
                else:
                    # if media_type == 'image/png':
                    #     return StreamingResponse(BytesIO(_cache[key]), media_type=media_type)
                    # else:
                    return StreamingResponse(BytesIO(_cache[key]), media_type="image/png")

            # Call the actual function if not in cache or expired
            response, image_data = await func(*args, **kwargs)
            # Cache the content of the response's body.
            _cache[key] = image_data
            _cache_time[key] = datetime.now()

            return response

        return wrapper

    return decorator