import asyncio
from typing import AsyncIterator

from grpclib.server import Server, Stream

from .proto.qlik.sse import BundledRows, Capabilities, ConnectorBase, Empty, FunctionRequestHeader
from .qlik_sse_function import QlikSSEFunction, QlikSSEFunctionWrapper, WrappableFunction


class QlikSSE(ConnectorBase):
    DEFAULT_HOST = "0.0.0.0"
    DEFAULT_PORT = 50055

    _name: str
    _version: str
    _functions: list[QlikSSEFunction]

    def __init__(self, name: str, version: str) -> None:
        self._name = name
        self._version = version
        self._functions = []

    async def get_capabilities(self, message: Empty) -> Capabilities:
        return Capabilities(
            allow_script=False,
            plugin_identifier=self._name,
            plugin_version=self._version,
            functions=[function.get_definition(idx) for idx, function in enumerate(self._functions)],
        )

    async def execute_function(self, messages: Stream[BundledRows, BundledRows]) -> AsyncIterator[BundledRows]:
        function_request_header = FunctionRequestHeader.parse(messages.metadata["qlik-functionrequestheader-bin"])
        async for bundle in self._functions[function_request_header.function_id](messages):
            yield bundle

    async def evaluate_script(self, messages: Stream[BundledRows, BundledRows]) -> AsyncIterator[BundledRows]:
        yield BundledRows()

    def function[F: WrappableFunction | QlikSSEFunction](self, function: F) -> F:
        if isinstance(function, QlikSSEFunction):
            self._functions.append(function)
        else:
            self._functions.append(QlikSSEFunctionWrapper(function))
        return function

    async def start(self, host: str = DEFAULT_HOST, port: int = DEFAULT_PORT) -> Server:
        server = Server([self])
        await server.start(host=host, port=port)
        return server

    async def serve(self, host: str = DEFAULT_HOST, port: int = DEFAULT_PORT) -> None:
        server = await self.start(host, port)
        await server.wait_closed()

    def serve_sync(self, host: str = DEFAULT_HOST, port: int = DEFAULT_PORT) -> None:
        asyncio.run(self.serve(host, port))
