import asyncio
import logging
import time
from typing import List, Dict, Any, Callable, Awaitable

from clarabot_nano.gateway.errors import SERVER_ERROR_NO_SESSION_BIND, SERVER_ERROR_TOO_MANY_REQUESTS
from clarabot_nano.gateway.state.gateway import GatewayState
from clarabot_nano.gateway.utility import check_single_response, GatewayServerError

logger = logging.getLogger(__name__)


async def state_session_job(
    gws: GatewayState, job_queue: asyncio.Queue, set_gateway_status: Callable[[str], Awaitable]
):
    """
    Manage the connection and session by handling events and publishing served rooms.
    Advertise/assert the served resources through the server if we acquire an active session. Periodically
    or when the config changes, refresh the served resources on the server.
    """
    serves = gws.extracted_config.get_serves()
    serve_assert_t = time.perf_counter() + 60
    assert_serve_now = False
    job_future = asyncio.ensure_future(job_queue.get())
    try:
        while True:
            done, _ = await asyncio.wait([job_future], timeout=serve_assert_t - time.perf_counter())
            if done:
                job = job_future.result()
                job_future = asyncio.ensure_future(job_queue.get())
                # handle job/event
                if (t := job.get("type")) == "session-change":
                    assert_serve_now = True
                    await set_gateway_status("active" if job.get("status") == "new" else "connected")
                elif t == "serve-change":  # Local config has changed.
                    serves = job["serves"]
                    assert_serve_now = True
            else:
                assert_serve_now = True
            # Assert what we want to serve, this will happen in response to events or the timeout.
            if assert_serve_now:
                try:
                    await _work_assert_serve(serves, gws)
                except GatewayServerError as e:
                    # Rate-limiting could knock the Nano offline temporarily, but it will try to serve regularly.
                    if e.error != SERVER_ERROR_TOO_MANY_REQUESTS:
                        raise
                    else:
                        serve_assert_t = time.perf_counter() + 5  # retry in a few seconds
                else:
                    serve_assert_t = time.perf_counter() + 210  # lease is ~5 minutes on server side
                assert_serve_now = False
    finally:
        job_future.cancel()
        await asyncio.gather(job_future, return_exceptions=True)


async def _work_assert_serve(serves: List[Dict[str, Any]], gws: GatewayState):
    """
    Tell the server which rooms are to be served by this Nano.
    Errors here are fatal.
    """
    try:
        async with gws.server_request({"type": "serve-rooms", "serves": serves}) as future:
            resp = check_single_response(await future, dict, "set served rooms")
    except GatewayServerError as e:
        # NOTE: Realistically no other errors should arrive here. The endpoint handles invalid values as soft-error.
        if e.error == SERVER_ERROR_NO_SESSION_BIND:
            room_ids = set()
        else:
            raise e
    else:
        if not isinstance(served_ids := resp.get("served_ids"), list) or not isinstance(
            invalid_ids := resp.get("invalid_ids"), list
        ):
            raise Exception(f"failed to set served rooms ({', '.join(resp.keys())})")
        if invalid_ids:
            logger.warning("Invalid room ids cannot be served: %s", invalid_ids)
        room_ids = set(serve["id"] for serve in serves)
        room_ids.intersection_update(s for s in served_ids if isinstance(s, int))
    gws.handle_serves_change(room_ids)
