from time import sleep
from concurrent.futures import ThreadPoolExecutor
from tracker import db, config

from api import models
from tracker import cacher

# DOCS https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ThreadPoolExecutor
class XRaySafeThreadPoolExecutor(ThreadPoolExecutor):
    def submit(self, *args, **kwargs):
        if config.enable_aws_xray:
            from aws_xray_sdk.core import xray_recorder
            xray_recorder.begin_segment('threadpoolexec')
        try:
            super().submit(*args, **kwargs)
        finally:
            xray_recorder.end_segment()


executor = ThreadPoolExecutor(2)


def cache_artists(source, ids):
    assert source in ('a', 'sc', 'in', 'tw', 'spy', 'yt'), f"Must be a valid source, was {source}"

    if config.enable_aws_xray:
        from aws_xray_sdk.core import xray_recorder
        try:
            xray_recorder.current_segment()
        except:
            xray_recorder.begin_segment('cache_artists')

    print(f"Caching Artists: {len(ids)} {source} / {ids}")
    if not ids:
        print(f'No ids to cache.')
        return

    print("setting db setup")
    db.setup_session()
    print("db setup")

    try:
        artists_to_cache = [a['artist'] for a in models.artists(source, ids, include_alerts=True, include_media=True, use_cache_only=False)]
    except Exception as e:
        print(f"Failed to fetch artists {e}")
        import traceback
        traceback.print_exc()


    if not artists_to_cache:
        print(f"Warning, failed to find artists")
    print(f"found {len(artists_to_cache)} artists to cache")
    cacher.set_artists(artists_to_cache)
    print(f"Cached {len(artists_to_cache)} artists.")


def test_long_task(arg1, arg2):
    print("Task #2 started with args: %s %s!" % (arg1, arg2))
    sleep(5)
    print("Task #2 is done!")