from __future__ import annotations
import asyncio
from typing import Any
from ._script import SSE_SCRIPT
from .broadcaster import Broadcaster
from .watcher import FileWatcher


class AutoreloadMiddleware:
    def __init__(
        self,
        app: Any,
        watch_paths: list[str] = ["."],
        sse_path: str = "/_autoreload",
        build_command: list[str] | None = None,
        exclude_paths: list[str] | None = None,
    ) -> None:
        self._app = app
        self._watch_paths = watch_paths
        self._sse_path = sse_path
        self._broadcaster = Broadcaster()
        self._watcher = FileWatcher(
            watch_paths, self._broadcaster, build_command, exclude_paths
        )
        self._watcher_started = False

    async def __call__(self, scope: dict, receive: Any, send: Any) -> None:
        if scope["type"] == "lifespan":
            await self._handle_lifespan(scope, receive, send)
        elif scope["type"] == "http" and scope["path"] == self._sse_path:
            await self._handle_sse(scope, receive, send)
        elif scope["type"] == "http":
            await self._handle_http(scope, receive, send)
        else:
            await self._app(scope, receive, send)

    async def _handle_lifespan(self, scope: dict, receive: Any, send: Any) -> None:
        async def wrapped_receive() -> dict:
            message = await receive()
            if message["type"] == "lifespan.startup":
                self._start_watcher()
            elif message["type"] == "lifespan.shutdown":
                await self._watcher.stop()
            return message

        await self._app(scope, wrapped_receive, send)

    def _start_watcher(self) -> None:
        if not self._watcher_started:
            self._watcher.start()
            self._watcher_started = True

    async def _handle_sse(self, scope: dict, receive: Any, send: Any) -> None:
        self._start_watcher()
        queue = await self._broadcaster.subscribe()

        await send({
            "type": "http.response.start",
            "status": 200,
            "headers": [
                [b"content-type", b"text/event-stream"],
                [b"cache-control", b"no-cache"],
                [b"connection", b"keep-alive"],
            ],
        })

        async def sse_task() -> None:
            while True:
                event = await queue.get()
                data = f"event: {event}\ndata: \n\n".encode()
                await send({"type": "http.response.body", "body": data, "more_body": True})

        async def disconnect_task() -> None:
            while True:
                msg = await receive()
                if msg["type"] == "http.disconnect":
                    return

        t1 = asyncio.ensure_future(sse_task())
        t2 = asyncio.ensure_future(disconnect_task())
        try:
            done, pending = await asyncio.wait(
                {t1, t2}, return_when=asyncio.FIRST_COMPLETED
            )
            for t in pending:
                t.cancel()
        finally:
            await self._broadcaster.unsubscribe(queue)
            await send({"type": "http.response.body", "body": b"", "more_body": False})

    async def _handle_http(self, scope: dict, receive: Any, send: Any) -> None:
        interceptor = _ResponseInterceptor(send)
        await self._app(scope, receive, interceptor)
        await interceptor.flush()


class _ResponseInterceptor:
    def __init__(self, send: Any) -> None:
        self._send = send
        self._start: dict | None = None
        self._body = b""

    async def __call__(self, message: dict) -> None:
        if message["type"] == "http.response.start":
            self._start = message
        elif message["type"] == "http.response.body":
            self._body += message.get("body", b"")
        else:
            await self._send(message)

    async def flush(self) -> None:
        if self._start is None:
            return

        headers: list[list[bytes]] = list(self._start.get("headers", []))
        content_type = b""
        content_encoding = b""
        for name, value in headers:
            if name.lower() == b"content-type":
                content_type = value
            if name.lower() == b"content-encoding":
                content_encoding = value

        body = self._body
        if b"text/html" in content_type and not content_encoding:
            if b"</body>" in body:
                body = body.replace(b"</body>", SSE_SCRIPT + b"</body>", 1)
            else:
                body = body + SSE_SCRIPT
            headers = [
                [n, v] for n, v in headers if n.lower() != b"content-length"
            ]
            headers.append([b"content-length", str(len(body)).encode()])

        await self._send({**self._start, "headers": headers})
        await self._send({"type": "http.response.body", "body": body, "more_body": False})
