import asyncio
import sys
import time
from contextlib import asynccontextmanager

import aiohttp


class _ControlApiStatusCollection:
    def __init__(self):
        self.requests_per_second = []
        self.indexer_work = []
        self.err_messages = []
        self.state = {}

    def collect(self, msg: dict):
        t = msg.pop("type")
        if t == "state":
            self.state.update(msg)
        elif t == "req_stat":
            self.requests_per_second.append(msg["req_per_s"])
        elif t == "idx_work":
            self.indexer_work.append(msg["work"])
        elif t == "err_message":
            self.err_messages.append(msg["msg"])
        else:
            raise Exception(f"{t} {msg}")

    def get(self):
        d = {
            "rps_sum": sum(self.requests_per_second),
            "rps_count": len(self.requests_per_second),
            "idx_works": self.indexer_work,
            "errors": self.err_messages,
            "state": self.state,
        }
        self.requests_per_second = []
        self.indexer_work = []
        self.err_messages = []
        self.state = dict(self.state)
        return d

    async def wait_state(self, kw_, timeout=30):
        t_ = time.perf_counter() + timeout
        while True:
            c_ = self.get()
            for k, v in kw_.items():
                if c_["state"].get(k) != v:
                    break
            else:
                return c_
            assert time.perf_counter() < t_, str(c_)
            await asyncio.sleep(0.1)


class ControlApiTestingClient:
    def __init__(self, control_api_path: str):
        if sys.platform.startswith("win"):
            conn = aiohttp.NamedPipeConnector(path=control_api_path)
        else:
            conn = aiohttp.UnixConnector(path=control_api_path)
        self.session = aiohttp.ClientSession(connector=conn)
        # status collection
        self._collection = None
        self._ws = None

    async def close(self):
        await self.session.close()

    @asynccontextmanager
    async def stream_status(self):
        self._collection = _ControlApiStatusCollection()
        async with self.session.ws_connect("unix://sock/status") as ws:
            self._ws = ws
            t = asyncio.create_task(self._status_collector(self._collection, self._ws))
            try:
                yield self._collection
            finally:
                await ws.close()
                await t

    async def _status_collector(self, c: _ControlApiStatusCollection, ws: aiohttp.ClientWebSocketResponse):
        async for msg in ws:
            msg: aiohttp.WSMessage
            if msg.type == aiohttp.WSMsgType.TEXT:
                c.collect(msg.json())
            else:
                await ws.close()
                raise Exception(str(msg))

    async def do_login(self, email, password, code=None):
        data = {"email": email, "password": password}
        if code:
            data["2faCode"] = code
        async with self.session.post("unix://sock/do-login", json=data) as resp:
            resp.raise_for_status()
            return await resp.json()

    async def do_reset(self):
        async with self.session.post("unix://sock/do-reset") as resp:
            resp.raise_for_status()
            return await resp.json()

    async def do_new_license(self, new_lic):
        data = {"new_license": new_lic}
        async with self.session.post("unix://sock/new-license", json=data) as resp:
            resp.raise_for_status()
            return await resp.json()

    async def do_rename(self, name):
        data = {"name": name}
        async with self.session.post("unix://sock/rename", json=data) as resp:
            resp.raise_for_status()
            return await resp.json()

    async def wait_gateway_active(self):
        gw_status = None
        for _ in range(60):
            await asyncio.sleep(0.5)
            c = self._collection.get()
            assert c.pop("rps_count") >= 0
            s = c.pop("state")
            gw_status = s.pop("gateway_status")
            assert s.get("indexer_status") == "ready"
            assert s.get("gateway_update") == 0
            assert c == {"rps_sum": 0, "idx_works": [], "errors": []}, str(c)
            if gw_status == "active":
                break
            assert gw_status == "connected", str(gw_status)
        else:
            assert False, str(gw_status)
