import asyncio
from typing import Optional, Any

import orjson
import bson
import aiohttp

from lacia.network.abcbase import BaseClient
from lacia.logger import logger
from lacia.types import Message
from lacia.exception import JsonRpcWsConnectException


class AioClient(BaseClient):
    def __init__(
        self,
        path: str = "",
        host: str = "localhost",
        port: int = 8080,
        loop: Optional[asyncio.AbstractEventLoop] = None,
        session: Optional[aiohttp.ClientSession] = None,
        headers: Optional[dict[str, Any]] = None,
        **kwargs,
    ) -> None:
        self.path = path
        self.host = host
        self.port = port
        self.loop = loop
        self.session = session
        self.kwargs = kwargs
        self.headers = headers

    def _get_headers(self, **kwargs):
        return {
            **self.kwargs,
            "headers": self.headers,
            **kwargs,
        }

    async def start(self, **kwargs) -> "AioClient":
        if self.session is None:
            self.session = aiohttp.ClientSession(
                loop=self.loop or asyncio.get_event_loop(),
                **self._get_headers(**kwargs),
            )
        self.ws = await self.session.ws_connect(
            f"http://{self.host}:{self.port}{self.path}"
        )
        logger.success(
            f"📡 {self.__class__.__name__} success connected: http://{self.host}:{self.port}{self.path}."
        )
        return self

    async def receive(self):
        try:
            async for data in self.ws:
                if data.type == aiohttp.WSMsgType.close:
                    raise JsonRpcWsConnectException(data.data)
                else:
                    return data
        except Exception as e:
            raise JsonRpcWsConnectException(f"{self.__class__.__name__} closed.({e})")

    async def receive_json(self):
        data = await self.receive()
        if data and data.type == aiohttp.WSMsgType.TEXT:
            data = orjson.loads(data.data)
            return data
        elif data and data.type == aiohttp.WSMsgType.BINARY:
            data = bson.loads(data.data)
            return data
        raise JsonRpcWsConnectException("Invalid data type.")

    async def receive_bytes(self):
        data = await self.receive()
        if data and data.type == aiohttp.WSMsgType.BINARY:
            return data.data

    async def iter_json(self):
        try:
            while True:
                data = await self.receive_json()
                if data:
                    yield data
        except JsonRpcWsConnectException as e:
            logger.error(f"http://{self.host}:{self.port}{self.path} closed.")
            await self.close()
        except Exception as e:
            logger.error(e)
            logger.error(f"http://{self.host}:{self.port}{self.path} closed.")
            await self.close()
        finally:
            return

    async def iter_bytes(self):
        try:
            while True:
                data = await self.receive_bytes()
                if data:
                    yield data
        except JsonRpcWsConnectException as e:
            logger.info(f"http://{self.host}:{self.port}{self.path} closed.")
            await self.close()
        except Exception as e:
            logger.error(e)
            logger.info(f"http://{self.host}:{self.port}{self.path} closed.")
            await self.close()

    async def send(self, message) -> None:
        return await self.send(message)

    async def send_bytes(self, message: bytes):
        return await self.ws.send_bytes(message)

    async def send_json(self, message: Message, binary: bool = True):
        if binary:
            return await self.ws.send_bytes(bson.dumps(message))
        return await self.ws.send_json(message)

    async def close(self) -> None:
        await self.ws.close()

    def closed(self) -> bool:
        return self.ws.closed

    def add_headers(self, **kwargs) -> None:
        if self.headers is None:
            self.headers = {}
        self.headers = {**self.headers, **kwargs}
