from __future__ import annotations

import asyncio
import contextlib
import pathlib
from functools import partial, reduce
from io import IOBase
from typing import (
    TYPE_CHECKING,
    Any,
    AsyncContextManager,
    Callable,
    Dict,
    Generic,
    List,
    Optional,
    TypeVar,
    Union,
    overload,
)
from weakref import WeakValueDictionary

import aiohttp
from aiohttp import (
    ClientResponse,
    ClientSession,
    ClientTimeout,
    ClientWebSocketResponse,
    WSMsgType,
    web,
)
from launart import ExportInterface, Service
from launart.manager import Launart
from launart.utilles import wait_fut
from loguru import logger
from statv import Stats
from typing_extensions import ParamSpec, Self

from graia.amnesia.json import Json, TJson
from graia.amnesia.transport import Transport
from graia.amnesia.transport.common.http import AbstractServerRequestIO
from graia.amnesia.transport.common.http import HttpEndpoint as HttpEndpoint
from graia.amnesia.transport.common.http.extra import HttpRequest, HttpResponse
from graia.amnesia.transport.common.http.io import AbstractClientRequestIO
from graia.amnesia.transport.common.status import ConnectionStatus
from graia.amnesia.transport.common.websocket import AbstractWebsocketIO
from graia.amnesia.transport.common.websocket import (
    WebsocketEndpoint as WebsocketEndpoint,
)
from graia.amnesia.transport.common.websocket.event import (
    WebsocketCloseEvent as WebsocketCloseEvent,
)
from graia.amnesia.transport.common.websocket.event import (
    WebsocketConnectEvent as WebsocketConnectEvent,
)
from graia.amnesia.transport.common.websocket.event import (
    WebsocketReceivedEvent as WebsocketReceivedEvent,
)
from graia.amnesia.transport.common.websocket.event import (
    WebsocketReconnect as WebsocketReconnect,
)
from graia.amnesia.transport.common.websocket.io import AbstractWebsocketIO
from graia.amnesia.transport.common.websocket.operator import (
    WSConnectionAccept as WebsocketAccept,
)
from graia.amnesia.transport.common.websocket.operator import (
    WSConnectionClose as WebsocketClose,
)
from graia.amnesia.transport.exceptions import ConnectionClosed
from graia.amnesia.transport.rider import TransportRider
from graia.amnesia.transport.signature import TransportSignature
from graia.amnesia.utilles import random_id


class AiohttpConnectionStatus(ConnectionStatus):
    drop = Stats[bool]("drop", default=False)

    def __init__(self) -> None:
        super().__init__()

    def __repr__(self) -> str:
        return f"<ConnectionStatus {self.connected=:}, {self.succeed=:}, {self.drop=:}, {self._waiters=:}>"

    async def wait_for_drop(self) -> None:
        while not self.drop:
            await self.wait_for_update()


class ClientRequestIO(AbstractClientRequestIO):
    rider: ClientConnectionRider[ClientResponse]
    response: ClientResponse

    def __init__(self, rider: ClientConnectionRider) -> None:
        assert rider.response
        self.rider = rider
        self.response = rider.response

    async def read(self) -> bytes:
        return await self.response.read()

    async def extra(self, signature):
        if signature is HttpResponse:
            return HttpResponse(
                self.response.status,
                dict(self.response.headers),
                {k: str(v) for k, v in self.response.cookies.items()},
                self.response.url,
            )

    def close(self):
        self.rider.status.drop = True


class ClientWebsocketIO(AbstractWebsocketIO):
    rider: ClientConnectionRider[ClientWebSocketResponse]
    connection: ClientWebSocketResponse

    def __init__(self, rider: ClientConnectionRider) -> None:
        assert rider.response
        self.rider = rider
        self.connection = rider.response

    async def cookies(self) -> Dict[str, str]:
        return {k: v.value for k, v in self.connection._response.cookies.items()}

    async def extra(self, signature):
        if signature is WebsocketClose:
            if not self.connection.closed:
                await self.connection.close()
        elif signature is HttpResponse:
            return HttpResponse(
                self.connection._response.status,
                dict(self.connection._response.request_info.headers),
                {k: v.value for k, v in self.connection._response.cookies.items()},
                self.connection._response.request_info.url,
            )

    async def receive(self) -> Union[bytes, str]:
        msg = await self.connection.receive()
        if msg.type in {WSMsgType.TEXT, WSMsgType.BINARY}:
            return msg.data
        # 错误处理
        if msg.type in {WSMsgType.CLOSE, WSMsgType.CLOSED}:
            raise ConnectionClosed("websocket closed")
        elif msg.type == WSMsgType.ERROR:
            raise msg.data
        else:
            raise TypeError(f"unexpected websocket message type: {msg.type}")

    async def send(self, data: "str | bytes | Any"):
        try:
            if isinstance(data, str):
                await self.connection.send_str(data)
            elif isinstance(data, bytes):
                await self.connection.send_bytes(data)
            else:
                await self.connection.send_json(data)
        except ConnectionError as e:
            raise ConnectionClosed(e.__class__.__qualname__, *e.args) from None

    @property
    def closed(self):
        return self.connection.closed

    async def wait_for_ready(self):
        if self.closed:
            raise ConnectionClosed("websocket closed")


T = TypeVar("T", ClientResponse, ClientWebSocketResponse)

P = ParamSpec("P")


class ClientConnectionRider(TransportRider[str, T], Generic[T]):
    def __init__(
        self,
        interface: AiohttpClientInterface,
        conn_func: Callable[..., AsyncContextManager[T]],
        call_param: Dict[str, Any],
    ) -> None:
        self.transports: List[Transport] = []
        self.connections = {}
        self.response: Optional[T] = None
        self.conn_func = conn_func
        self.call_param = call_param
        self.status = AiohttpConnectionStatus()
        self.autoreceive: bool = False
        self.task = None
        self.connect_task = None
        self.interface = interface

    async def _connect(self):
        async with self.conn_func(**self.call_param) as resp:
            self.response = resp
            self.status.connected = True
            await self.status.wait_for_drop()
            self.status.drop = False
        self.status.connected = False

    async def _start_conn(self) -> Self:
        if not self.status.available:
            self.response = None
            if self.connect_task is None:
                self.connect_task = asyncio.create_task(self._connect())
            await wait_fut((self.status.wait_for_available(), self.connect_task), return_when=asyncio.FIRST_COMPLETED)
            if self.connect_task.done():
                exc = self.connect_task.exception()
                self.connect_task = None
                if exc:
                    raise exc
            self.status.succeed = True
        return self

    def __await__(self):
        return self._start_conn().__await__()

    @overload
    def io(self: "ClientConnectionRider[ClientResponse]") -> ClientRequestIO:
        ...

    @overload
    def io(self: "ClientConnectionRider[ClientWebSocketResponse]") -> ClientWebsocketIO:
        ...

    def io(self, id=None) -> ...:
        if id:
            raise TypeError("this rider has just one connection")
        if self.autoreceive:
            raise TypeError("this rider has been taken over by auto receive, use .use(transport) instead.")
        if not self.status.connected:
            raise RuntimeError("the connection is not ready, please await the instance to ensure connection")
        assert self.response
        if isinstance(self.response, ClientWebSocketResponse):
            return ClientWebsocketIO(self)
        elif isinstance(self.response, ClientResponse):
            return ClientRequestIO(self)
        else:
            raise TypeError("this response is not a ClientResponse or ClientWebSocketResponse")

    async def trigger_callbacks(self, event: TransportSignature[Callable[P, Any]], *args: P.args, **kwargs: P.kwargs):
        tasks: List[asyncio.Task] = []
        for i in self.transports:
            if i.has_callback(event):
                tasks.extend(asyncio.create_task(f(*args, **kwargs)) for f in i.get_callbacks(event))
        for task in tasks:
            await task

    async def connection_manage(self):
        __original_transports: List[Transport] = self.transports[:]

        with contextlib.suppress(Exception):
            while self.transports and self.interface.service.status.stage != "finished":
                try:
                    await self._start_conn()
                    assert isinstance(
                        self.response, ClientWebSocketResponse
                    ), f"{self.response} is not a ClientWebSocketResponse"
                    io = ClientWebsocketIO(self)
                    await self.trigger_callbacks(WebsocketConnectEvent, io)
                    with contextlib.suppress(ConnectionClosed):
                        async for data in io.packets():
                            await self.trigger_callbacks(WebsocketReceivedEvent, io, data)
                    await self.trigger_callbacks(WebsocketCloseEvent, io)
                    if not io.closed:
                        await io.close()
                    self.status.drop = True
                    await self.status.wait_for_unavailable()
                    self.connect_task = None
                except aiohttp.ClientConnectionError as e:
                    logger.warning(repr(e))
                except Exception as e:
                    logger.exception(e)
                # scan transports
                if self.interface.service.status == "cleanup":
                    continue
                continuing_transports: List[Transport] = self.transports[:]
                self.transports = []
                reconnect_handle_tasks = []
                for t in continuing_transports:
                    if t.has_handler(WebsocketReconnect):
                        handler = t.get_handler(WebsocketReconnect)
                        tsk = asyncio.create_task(handler(self.status))
                        tsk.add_done_callback(lambda tsk: self.transports.append(t) if tsk.result() is True else None)
                        reconnect_handle_tasks.append(tsk)
                await wait_fut(reconnect_handle_tasks)
        self.transports = __original_transports

    def use(self, transport: Transport):
        self.autoreceive = True
        self.transports.append(transport)
        if not self.task:
            self.task = asyncio.create_task(self.connection_manage())
        return self.task


class AiohttpClientInterface(ExportInterface["AiohttpService"]):
    service: AiohttpService

    def __init__(self, service: AiohttpService) -> None:
        self.service = service

    def request(
        self,
        method: str,
        url: str,
        params: Optional[dict] = None,
        data: Optional[Any] = None,
        headers: Optional[dict] = None,
        cookies: Optional[dict] = None,
        timeout: Optional[float] = None,
        *,
        json: Optional[TJson] = None,
        **kwargs: Any,
    ) -> ClientConnectionRider[ClientResponse]:
        if json:
            data = Json.serialize(json)
        call_param: Dict[str, Any] = {
            "method": method,
            "url": url,
            "params": params,
            "data": data,
            "headers": headers,
            "cookies": cookies,
            "timeout": timeout,
            **kwargs,
        }
        return ClientConnectionRider(self, self.service.session.request, call_param)

    def websocket(self, url: str, **kwargs) -> ClientConnectionRider[ClientWebSocketResponse]:
        call_param: Dict[str, Any] = {"url": url, **kwargs}
        return ClientConnectionRider(self, self.service.session.ws_connect, call_param)


class AiohttpService(Service):
    id = "http.universal_client"
    session: ClientSession
    supported_interface_types = {AiohttpClientInterface}

    def __init__(self, session: Optional[ClientSession] = None) -> None:
        if TYPE_CHECKING:
            session = session or ClientSession()
        self.session = session
        super().__init__()

    def get_interface(self, interface_type):
        if interface_type is AiohttpClientInterface:
            return AiohttpClientInterface(self)

    @property
    def stages(self):
        return {"preparing", "cleanup"}

    @property
    def required(self):
        return set()

    async def launch(self, mgr: Launart):
        async with self.stage("preparing"):
            if not self.session:
                self.session = ClientSession(timeout=ClientTimeout(total=None))
        async with self.stage("cleanup"):
            await self.session.close()


class AiohttpServerRequestIO(AbstractServerRequestIO):
    request: web.Request

    def __init__(self, request: web.Request) -> None:
        self.request = request

    async def read(self) -> bytes:
        return await self.request.content.read()

    async def extra(self, signature):
        if signature is HttpRequest:
            return HttpRequest(
                dict(self.request.headers),
                dict(self.request.cookies),
                dict(self.request.query),
                self.request.url,
                self.request.host,
                self.request.method,
                self.request.remote or ".0.0.0",
                0,
            )


class AiohttpServerWebsocketIO(AbstractWebsocketIO):
    websocket: web.WebSocketResponse
    request: web.Request
    ready: asyncio.Event

    def __init__(self, request: web.Request) -> None:
        self.websocket = web.WebSocketResponse(autoping=True)
        self.request = request
        self.ready = asyncio.Event()

    async def receive(self) -> Union[str, bytes]:
        try:
            received = await self.websocket.receive()
        except asyncio.CancelledError as e:
            raise ConnectionClosed("Cancelled") from e
        if received.type in (web.WSMsgType.BINARY, web.WSMsgType.TEXT):
            return received.data
        elif received.type in (web.WSMsgType.CLOSE, web.WSMsgType.CLOSING, web.WSMsgType.CLOSED):
            self.ready.clear()
            raise ConnectionClosed("Connection closed")
        elif received.type is web.WSMsgType.ERROR:
            exc = self.websocket.exception()
            raise ConnectionClosed("Websocket Error") from exc
        raise TypeError(f"Unknown type of received message {received}")

    async def send(self, data: Union[str, bytes, Any]):
        if isinstance(data, str):
            await self.websocket.send_str(data)
        elif isinstance(data, bytes):
            await self.websocket.send_bytes(data)
        else:
            await self.websocket.send_json(data)

    async def extra(self, signature):
        if signature is HttpRequest:
            return HttpRequest(
                dict(self.request.headers),
                dict(self.request.cookies),
                dict(self.request.query),
                self.request.url,
                self.request.host,
                self.request.method,
                self.request.remote or ".0.0.0",
                0,
            )
        elif signature is WebsocketAccept:
            await self.websocket.prepare(self.request)
            self.ready.set()
        elif signature is WebsocketClose:
            await self.websocket.close()

    async def wait_for_ready(self):
        return await self.ready.wait()

    async def accept(self):
        return await super().accept()

    async def headers(self) -> Dict[str, str]:
        return dict(self.websocket.headers)

    async def cookies(self) -> Dict[str, str]:
        return dict(self.request.cookies)

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

    @property
    def closed(self):
        return self.websocket.closed


class AiohttpRouter(
    ExportInterface["AiohttpServerService"],
    TransportRider[str, Union[AiohttpServerRequestIO, AiohttpServerWebsocketIO]],
):
    def __init__(self, wsgi: web.Application) -> None:
        self.connections = WeakValueDictionary()
        self.transports = []
        self.wsgi = wsgi

    async def handle_http_request(self, handler: Callable[[AiohttpServerRequestIO], Any], request: web.Request):
        req_io = AiohttpServerRequestIO(request)
        conn_id = random_id()
        self.connections[conn_id] = req_io
        try:
            handler_resp = await handler(req_io)
        finally:
            del self.connections[conn_id]
        if not isinstance(handler_resp, tuple):
            body = handler_resp
            resp_info = {}
        else:
            body = handler_resp[0]
            resp_info = reduce(lambda x, y: dict(x, **y), handler_resp[1:], {})
        status: int = resp_info.get("status", 200)
        if isinstance(body, (str, bytes, IOBase)):
            response = web.Response(body=body, status=status)
        elif isinstance(body, (dict, list)):
            response = web.json_response(body, status=status, headers=resp_info.get("headers"))
        elif isinstance(body, pathlib.Path):
            response = web.Response(body=body.open("rb"), status=status)
        elif isinstance(body, web.Response):
            response = body
        else:
            raise ValueError(f"unsupported response type {type(body)}")

        if "cookies" in resp_info:
            expire = resp_info.get("cookie_expires")
            for key, value in resp_info["cookies"].items():
                response.set_cookie(key, value, expires=expire)

        return response

    async def handle_websocket_request(self, request: web.Request) -> web.WebSocketResponse:
        websocket_io = AiohttpServerWebsocketIO(request)
        conn_id = random_id()
        self.connections[conn_id] = websocket_io
        await self.trigger_callbacks(WebsocketConnectEvent, websocket_io)
        if websocket_io.closed:
            return websocket_io.websocket
        with contextlib.suppress(ConnectionClosed):
            async for message in websocket_io.packets():
                await self.trigger_callbacks(WebsocketReceivedEvent, websocket_io, message)
        await self.trigger_callbacks(WebsocketCloseEvent, websocket_io)
        await websocket_io.close()
        return websocket_io.websocket

    async def trigger_callbacks(self, event, *args, **kwargs):
        callbacks = [i.get_callbacks(event) for i in self.transports if i.has_callback(event)]
        if callbacks:
            callbacks = reduce(lambda a, b: a + b, callbacks)
            await wait_fut([i(*args, **kwargs) for i in callbacks])

    def io(self, id: Optional[str] = None):
        if id is None:
            raise ValueError("id is required")

        return self.connections.get(id)

    def use(self, transport: Transport):
        self.transports.append(transport)
        for signature, handler in transport.iter_handlers():
            if isinstance(signature, HttpEndpoint):
                for method in signature.methods:
                    self.wsgi.router.add_route(method, signature.path, partial(self.handle_http_request, handler))

        for signature in transport.declares:
            if isinstance(signature, WebsocketEndpoint):
                self.wsgi.router.add_route("GET", signature.path, self.handle_websocket_request)


class AiohttpServerService(Service):
    id = "http.universal_server"
    wsgi_handler: web.Application
    supported_interface_types = {AiohttpRouter}

    def __init__(
        self, host: str = "127.0.0.1", port: int = 8000, wsgi_handler: Optional[web.Application] = None
    ) -> None:
        self.wsgi_handler = wsgi_handler or web.Application()
        self.wsgi_handler.router.freeze = lambda: None  # monkey patch
        self.routers: List[AiohttpRouter] = []
        self.host = host
        self.port = port
        super().__init__()

    def get_interface(self, interface_type):
        if interface_type is AiohttpRouter:
            router = AiohttpRouter(self.wsgi_handler)
            self.routers.append(router)
            return router

    @property
    def stages(self):
        return {"preparing", "blocking", "cleanup"}

    @property
    def required(self):
        return set()

    async def launch(self, manager: Launart):
        async with self.stage("preparing"):
            logger.info(f"starting server on {self.host}:{self.port}")
            runner = web.AppRunner(self.wsgi_handler)
            await runner.setup()
            site = web.TCPSite(runner, self.host, self.port)
        async with self.stage("blocking"):
            await site.start()
        async with self.stage("cleanup"):
            await self.wsgi_handler.shutdown()
            await self.wsgi_handler.cleanup()
