"""
 * noload
"""
import os
import json
import asyncio
import toml
import logging
import traceback
from typing import Dict, List, NoReturn, Awaitable, Iterator, Tuple, Any, Optional
from asyncio import Future, Task, FIRST_COMPLETED
from datetime import timedelta

logger = logging.getLogger('server-query')

def kv_parse(string: str) -> Dict[str, str]:
    data = {}
    for kv in string.split(' '):
        try:
            val = kv.split('=', 1)[1]
        except IndexError:
            val = ''
        data[kv.split('=')[0]] = val.replace(r'\s', ' ')
    return data

class SQResult:
    def __init__(self, data: str) -> None:
        self.data = data

    def kvs(self) -> Dict[str, str]:
        return kv_parse(self.data)

    def split(self, token: str) -> List['SQResult']:
        return [SQResult(part) for part in self.data.split(token)]

class SQError(Exception):
    def __init__(self, message: str, errno: int):
        super().__init__(f'{errno}: {message}')
        self.errno = errno

class ServerQueryClient:
    def __init__(self, username: str, password: str,
                 host: str = '127.0.0.1', port: int = 10011) -> None:
        self.running = True
        self.host = host
        self.port = port
        self.username = username
        self.password = password

        self.request_queue: asyncio.Queue = asyncio.Queue()
        self.notify_queue: asyncio.Queue = asyncio.Queue()

        self.task = asyncio.create_task(self.async_client())
        self.ping_task = asyncio.create_task(self.ping())
        self.loop = asyncio.get_event_loop()
        self.debug = True

    @classmethod
    async def create(cls, *args, **kwargs) -> 'ServerQueryClient': # type: ignore
        self = cls(*args, **kwargs)
        await self.raw_request(f'login {self.username} {self.password}')
        await self.raw_request(f'use 1')
        return self

    async def ping(self) -> None:
        while self.running:
            await asyncio.sleep(60)
            response = await self.raw_request('version')

    def raw_request(self, req: str) -> asyncio.Future:
        fut = self.loop.create_future()
        # The queue has no limit, so put_nowait works.
        # This allows us to just return the future, which means we won't
        # get odd syntax like await await self.raw_request(...)
        self.request_queue.put_nowait((req, fut))
        return fut

    def request(self, req: str) -> asyncio.Future:
        fut = self.raw_request(req)
        new_fut = self.loop.create_future()
        def _callback(fut: asyncio.Future) -> None:
            try:
                result = fut.result()
                new_fut.set_result(SQResult(result))
            except SQError as e:
                new_fut.set_exception(e)
        fut.add_done_callback(_callback)
        return new_fut

    async def stop(self) -> None:
        self.running = False
        self.task.cancel()
        self.ping_task.cancel()

    def __await__(self) -> Iterator[Tuple[Any, Any]]:
        return iter(asyncio.gather(self.task, self.ping_task))

    def get_error(self, in_data: str) -> Tuple[str, int]:
        try:
            errstr = in_data.split('\n')[-2].strip()
        except IndexError:
            errstr = in_data.strip()
        assert errstr.startswith('error')

        data = kv_parse(errstr)
        return data['msg'].replace(r'\s', ' '), int(data['id'])

    async def process_notification(self, line: str) -> None:
        assert line.startswith('notify')
        logger.info(f'Received notification: {line}')
        res = SQResult(line)
        await self.notify_queue.put(res)

    async def readline(self) -> str:
        return (await self.reader.readline()).strip().decode('utf-8')

    async def handle_req_or_notify(self, req_aw: Optional[Task] = None, notify_aw: Optional[Task] = None) -> Tuple[Optional[Task], Optional[Task]]:
        if req_aw is None:
            logger.info('Starting new request queue task.')
            req_aw = asyncio.create_task(self.request_queue.get(), name='request task')
        if notify_aw is None:
            logger.info('Starting new notify (network readuntil) task.')
            notify_aw = asyncio.create_task(self.readline(), name='notify task')

        done, pending = await asyncio.wait({req_aw, notify_aw},
                                           return_when=FIRST_COMPLETED)

        if notify_aw in done:
            line = await notify_aw
            logger.info(f'Received notification line: {line}')
            await self.process_notification(line)
            logger.info(f'Processed notification line. notify_aw done, returning old req_aw.')
            return (req_aw, None)

        request, future = req_aw.result()

        self.writer.write((request + '\n').encode('utf-8'))
        await self.writer.drain()

        logger.info('Starting request reading...')

        req_data = ''
        line = await notify_aw # the 'notify' task will contain the first response str
        logger.info(f'First line read: {line!r}')

        if not 'error' in line:
            req_data += line
            while not (line := (await self.readline())).startswith('error'):
                if line.startswith('notify'):
                    await self.process_notification(line)

                req_data += line + '\n'

        error_data = kv_parse(line)
        logger.info(f'Req: {request}. Read response: {req_data.strip()}. Error line: {line.strip()}')

        if error_data['id'] != '0':
            future.set_exception(SQError(error_data['msg'], int(error_data['id'])))
        else:
            future.set_result(req_data.strip())

        return (None, None)

    async def async_client(self) -> None:
        self.reader, self.writer = await asyncio.open_connection(self.host, self.port)
        welcome_message = await self.reader.read(16384)

        logger.info("[sq_client] Read welcome message.")
        args: Tuple[Optional[Task], Optional[Task]] = (None, None)

        try:
            while self.running:
                args = await self.handle_req_or_notify(*args)
        finally:
            self.writer.close()
            await self.writer.wait_closed()