from collections import deque

from async_task_manager import ThreadWorkerQueue, io_task
from aws_srp import AWSSRP
from subscribe import Subscription
import json
import os
from id_test_devel import *
from time import sleep
from appsync_query import appsync_query_iam
from datetime import datetime
import logging
from copy import deepcopy

def make_iso_date_now():
    return f'{datetime.utcnow().isoformat(timespec="microseconds")}Z'


logging.basicConfig(filename=f'integrity_{make_iso_date_now()}.log',
                    level=logging.ERROR,
                    format='%(asctime)s %(message)s'
                    )
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

THREAD_COUNT = os.getenv("ASYNC_WORKERS", 2)
IO_THREAD_COUNT = os.getenv("IO_WORKERS", 100)

ThreadWorkerQueue.instantiate(thread_count=THREAD_COUNT, io_worker_limit=IO_THREAD_COUNT)

appsync_channel_id = 'a9028389dfd904ce51a2aaad93299463286b9ae90c3ee12ec16ea8d3cd37d99e'

QUERY_listenOnCollectionEvents_with_params = json.dumps({
    'query': """subscription listen($channelId:ID!) {
  listenOnCollectionEvents(channelId:$channelId){
    channelId
    collectionId
    collection {
        id
        name
        status
        parentId
        dateCreated
        type
        totalProfiles
    }
  }
}""",
    'variables': {"channelId": appsync_channel_id}
})

QUERY_send_collection = """mutation updateCollection($channelId:ID!, $collectionId: ID!, $collection:CollectionInput!) {
    sendCollection(channelId:$channelId, collectionId:$collectionId, collection:$collection)
    {
      channelId
      collectionId
      collection {
        id
        status
      }
      }
    }"""

send_variables = {'channelId': appsync_channel_id,
                  'collectionId': 1,
                  'collection': {'id': 1,
                                 'status': 'uploaded'}
                  }

TEST_RUNNING = True
SENT_MESSAGES = deque()
SUB = None

def listen_on_message(msg_type, msg_obj):
    log.info(f"ARRIVED: type: {msg_type}, obj: {msg_obj}")
    if msg_type == "data":
        try:
            got_id = int(msg_obj['payload']['data']['listenOnCollectionEvents']['collectionId'])
            SENT_MESSAGES.remove(got_id)
            if len(SENT_MESSAGES) > 0:
                log.info(f"SENT QUEUE: {', '.join([str(q) for q in SENT_MESSAGES])}")
        except Exception as e:
            log.exception("on_message: ", exc_info=e)


def listen_on_error(msg_obj):
    log.info(f"type: socket error, obj: {msg_obj}")


@io_task
def listener(task_handle=None):
    global SUB
    try:
        awssrp = AWSSRP(user_name, user_passwd, user_pool_id, app_id, pool_region)

        sub = Subscription(WSS_URL, API_URL, HOST, QUERY_listenOnCollectionEvents_with_params)

        SUB = sub
        while TEST_RUNNING:
            tokens = awssrp.authenticate_user()
            auth_token = tokens["AuthenticationResult"]["AccessToken"]
            # print(tokens)
            log.info(f"SUBSCRIBED to channel {appsync_channel_id} with id {sub.SUB_ID}")
            sub.sub(trace=False, on_message=listen_on_message, auth_token=auth_token, run_forever=False)
            task_handle.send_immediate_response("listening")
            sub.run_forever()
            log.info("LISTENER CLOSED")
    except Exception as e:
        task_handle.send_immediate_error(f"Oops: {e}")
        SUB = None
        raise e

@io_task
def as_send(send_vars, **kwargs):
    response = appsync_query_iam(API_URL, pool_region, QUERY_send_collection, send_vars)
    log.info(f"SENT: {response.content.decode()}")


@io_task
def sender(task_handle=None):
    n = 0
    while TEST_RUNNING:
        send_variables['collectionId'] = n
        send_variables['collection']['id'] = n
        SENT_MESSAGES.append(n)
        as_send(deepcopy(send_variables))
        n += 1
        if n % 100 > 10:
            sleep(1)
    return n


listen = listener()
if listen.wait_for_immediate_response() == 'listening':
    sleep(3)
    log.info("STARTING SENDER")
    send = sender()

    try:
        listen.wait_for_result()
    except Exception as e:
        print(f"CRAP? {e}")
    except KeyboardInterrupt as e:
        log.info("KEYBOARD INTERRUPT")
    finally:
        TEST_RUNNING = False
        if SUB:
            SUB.stop()
        sent_messages = send.wait_for_result()
        log.info(f"Sent messages: {sent_messages}")
        log.info(f"SENT QUEUE: {', '.join([str(q) for q in SENT_MESSAGES])}")
        ThreadWorkerQueue.instance.stop()