from __future__ import annotations

import asyncio
import json
import logging
import math
import struct
from typing import Any
from urllib.parse import urlparse

import click
import httpx
import websockets

# Bump this when making breaking changes to the WebSocket protocol.
# The server will reject clients with a version lower than its minimum.
PROTOCOL_VERSION = 2


class TunnelClient:
    def __init__(
        self, *, destination_url: str, subdomain: str, tunnel_host: str, log_level: str
    ) -> None:
        self.destination_url = destination_url
        self.subdomain = subdomain
        self.tunnel_host = tunnel_host

        if "localhost" in tunnel_host or "127.0.0.1" in tunnel_host:
            self.tunnel_http_url = f"http://{subdomain}.{tunnel_host}"
            self.tunnel_websocket_url = (
                f"ws://{subdomain}.{tunnel_host}?v={PROTOCOL_VERSION}"
            )
        else:
            self.tunnel_http_url = f"https://{subdomain}.{tunnel_host}"
            self.tunnel_websocket_url = (
                f"wss://{subdomain}.{tunnel_host}?v={PROTOCOL_VERSION}"
            )

        self.logger = logging.getLogger(__name__)
        level = getattr(logging, log_level.upper())
        self.logger.setLevel(level)
        self.logger.propagate = False
        handler = logging.StreamHandler()
        handler.setLevel(level)
        handler.setFormatter(logging.Formatter("%(message)s"))
        self.logger.addHandler(handler)

        self.pending_requests: dict[str, dict[str, Any]] = {}
        self.active_streams: dict[str, asyncio.Event] = {}
        self.stop_event = asyncio.Event()

    async def connect(self) -> None:
        retry_delay = 1.0
        max_retry_delay = 30.0
        while not self.stop_event.is_set():
            try:
                self.logger.debug(
                    f"Connecting to WebSocket URL: {self.tunnel_websocket_url}"
                )
                async with websockets.connect(
                    self.tunnel_websocket_url, max_size=None
                ) as websocket:
                    self.logger.debug("WebSocket connection established")
                    click.secho(
                        f"Connected to tunnel {self.tunnel_http_url}", fg="green"
                    )
                    retry_delay = 1.0
                    await self.handle_messages(websocket)
            except asyncio.CancelledError:
                self.logger.debug("Connection cancelled")
                break
            except websockets.InvalidStatus as e:
                if e.response.status_code == 426:
                    body = e.response.body.decode() if e.response.body else ""
                    click.secho(
                        body or "Client version too old. Please upgrade plain.tunnel.",
                        fg="red",
                    )
                    break
                raise
            except (websockets.ConnectionClosed, ConnectionError) as e:
                if self.stop_event.is_set():
                    self.logger.debug("Stopping reconnect attempts due to shutdown")
                    break
                click.secho(
                    f"Connection lost: {e}. Retrying in {retry_delay:.0f}s...",
                    fg="yellow",
                )
                await asyncio.sleep(retry_delay)
                retry_delay = min(retry_delay * 2, max_retry_delay)
            except Exception as e:
                if self.stop_event.is_set():
                    self.logger.debug("Stopping reconnect attempts due to shutdown")
                    break
                click.secho(
                    f"Unexpected error: {e}. Retrying in {retry_delay:.0f}s...",
                    fg="yellow",
                )
                await asyncio.sleep(retry_delay)
                retry_delay = min(retry_delay * 2, max_retry_delay)

    async def handle_messages(self, websocket: Any) -> None:
        try:
            async for message in websocket:
                if isinstance(message, str):
                    data = json.loads(message)
                    msg_type = data.get("type")
                    if msg_type == "ping":
                        self.logger.debug("Received heartbeat ping, sending pong")
                        await websocket.send(json.dumps({"type": "pong"}))
                    elif msg_type == "request":
                        self.logger.debug("Received request metadata from worker")
                        await self.handle_request_metadata(websocket, data)
                    elif msg_type == "stream-cancel":
                        request_id = data.get("id")
                        self.logger.debug(
                            f"Received stream-cancel for request ID: {request_id}"
                        )
                        cancel_event = self.active_streams.get(request_id)
                        if cancel_event:
                            cancel_event.set()
                    else:
                        self.logger.warning(
                            f"Received unknown message type: {msg_type}"
                        )
                elif isinstance(message, bytes):
                    self.logger.debug("Received binary data from worker")
                    await self.handle_request_body_chunk(websocket, message)
                else:
                    self.logger.warning("Received unknown message format")
        except asyncio.CancelledError:
            self.logger.debug("Message handling cancelled")
        except Exception as e:
            self.logger.error(f"Error in handle_messages: {e}")
            raise

    async def handle_request_metadata(
        self, websocket: Any, data: dict[str, Any]
    ) -> None:
        request_id = data["id"]
        has_body = data.get("has_body", False)
        total_body_chunks = data.get("totalBodyChunks", 0)
        self.pending_requests[request_id] = {
            "metadata": data,
            "body_chunks": {},
            "has_body": has_body,
            "total_body_chunks": total_body_chunks,
        }
        self.logger.debug(
            f"Stored metadata for request ID: {request_id}, has_body: {has_body}"
        )
        await self.check_and_process_request(websocket, request_id)

    async def handle_request_body_chunk(
        self, websocket: Any, chunk_data: bytes
    ) -> None:
        (id_length,) = struct.unpack_from("<I", chunk_data, 0)
        request_id = chunk_data[4 : 4 + id_length].decode("utf-8")
        header_end = 4 + id_length + 8
        chunk_index, total_chunks = struct.unpack_from("<II", chunk_data, 4 + id_length)
        body_chunk = chunk_data[header_end:]

        if request_id in self.pending_requests:
            request = self.pending_requests[request_id]
            request["body_chunks"][chunk_index] = body_chunk
            self.logger.debug(
                f"Stored body chunk {chunk_index + 1}/{total_chunks} for request ID: {request_id}"
            )
            await self.check_and_process_request(websocket, request_id)
        else:
            self.logger.warning(
                f"Received body chunk for unknown or completed request ID: {request_id}"
            )

    async def check_and_process_request(self, websocket: Any, request_id: str) -> None:
        request_data = self.pending_requests.get(request_id)
        if not request_data:
            return

        has_body = request_data["has_body"]
        total_body_chunks = request_data["total_body_chunks"]
        body_chunks = request_data["body_chunks"]

        all_chunks_received = not has_body or len(body_chunks) == total_body_chunks
        if not all_chunks_received:
            return

        for i in range(total_body_chunks):
            if i not in body_chunks:
                self.logger.error(
                    f"Missing chunk {i + 1}/{total_body_chunks} for request ID: {request_id}"
                )
                return

        self.logger.debug(f"Processing request ID: {request_id}")
        del self.pending_requests[request_id]
        task = asyncio.create_task(
            self.process_request(
                websocket,
                request_data["metadata"],
                body_chunks,
                request_id,
            )
        )
        task.add_done_callback(self._handle_task_exception)

    def _handle_task_exception(self, task: asyncio.Task[None]) -> None:
        if not task.cancelled() and task.exception():
            self.logger.error(f"Error processing request: {task.exception()}")

    async def process_request(
        self,
        websocket: Any,
        request_metadata: dict[str, Any],
        body_chunks: dict[int, bytes],
        request_id: str,
    ) -> None:
        self.logger.debug(
            f"Processing request: {request_id} {request_metadata['method']} {request_metadata['url']}"
        )

        if request_metadata["has_body"]:
            total_chunks = request_metadata["totalBodyChunks"]
            body_data = b"".join(body_chunks[i] for i in range(total_chunks))
        else:
            body_data = None

        parsed = urlparse(request_metadata["url"])
        path = parsed.path
        if parsed.query:
            path = f"{path}?{parsed.query}"
        forward_url = f"{self.destination_url}{path}"

        self.logger.debug(f"Forwarding request to: {forward_url}")

        async with httpx.AsyncClient(follow_redirects=False, verify=False) as client:
            try:
                async with client.stream(
                    method=request_metadata["method"],
                    url=forward_url,
                    headers=request_metadata["headers"],
                    content=body_data,
                ) as response:
                    response_status = response.status_code
                    response_headers = dict(response.headers)

                    self.logger.info(
                        f"{click.style(request_metadata['method'], bold=True)} {request_metadata['url']} {response_status}"
                    )

                    if self._is_streaming_response(response):
                        await self._handle_streaming_response(
                            websocket,
                            response,
                            request_id,
                            response_status,
                            response_headers,
                        )
                    else:
                        await response.aread()
                        await self._handle_buffered_response(
                            websocket,
                            response.content,
                            request_id,
                            response_status,
                            response_headers,
                        )
            except httpx.ConnectError as e:
                self.logger.error(f"Connection error forwarding request: {e}")
                self.logger.info(
                    f"{click.style(request_metadata['method'], bold=True)} {request_metadata['url']} 502"
                )
                await self._handle_buffered_response(
                    websocket, b"", request_id, 502, {}
                )

    def _is_streaming_response(self, response: httpx.Response) -> bool:
        content_type = response.headers.get("content-type", "")
        return "text/event-stream" in content_type

    async def _handle_buffered_response(
        self,
        websocket: Any,
        response_body: bytes,
        request_id: str,
        response_status: int,
        response_headers: dict[str, str],
    ) -> None:
        has_body = len(response_body) > 0
        max_chunk_size = 1_000_000
        total_body_chunks = (
            math.ceil(len(response_body) / max_chunk_size) if has_body else 0
        )

        response_metadata = {
            "type": "response",
            "id": request_id,
            "status": response_status,
            "headers": list(response_headers.items()),
            "has_body": has_body,
            "totalBodyChunks": total_body_chunks,
        }

        self.logger.debug(
            f"Sending response metadata for ID: {request_id}, has_body: {has_body}"
        )
        await websocket.send(json.dumps(response_metadata))

        if has_body:
            self.logger.debug(
                f"Sending {total_body_chunks} body chunks for ID: {request_id}"
            )
            id_bytes = request_id.encode("utf-8")
            for i in range(total_body_chunks):
                chunk_start = i * max_chunk_size
                chunk_end = min(chunk_start + max_chunk_size, len(response_body))
                header = id_bytes + struct.pack("<II", i, total_body_chunks)
                await websocket.send(header + response_body[chunk_start:chunk_end])
                self.logger.debug(
                    f"Sent body chunk {i + 1}/{total_body_chunks} for ID: {request_id}"
                )

    async def _handle_streaming_response(
        self,
        websocket: Any,
        response: httpx.Response,
        request_id: str,
        response_status: int,
        response_headers: dict[str, str],
    ) -> None:
        cancel_event = asyncio.Event()
        self.active_streams[request_id] = cancel_event

        stream_start = {
            "type": "stream-start",
            "id": request_id,
            "status": response_status,
            "headers": list(response_headers.items()),
        }

        self.logger.debug(f"Sending stream-start for ID: {request_id}")
        await websocket.send(json.dumps(stream_start))

        id_bytes = request_id.encode("utf-8")

        try:
            async for chunk in response.aiter_bytes():
                if cancel_event.is_set():
                    self.logger.debug(
                        f"Stream cancelled by browser for request ID: {request_id}"
                    )
                    break

                await websocket.send(id_bytes + chunk)

            stream_end = {
                "type": "stream-end",
                "id": request_id,
            }
            self.logger.debug(f"Sending stream-end for ID: {request_id}")
            await websocket.send(json.dumps(stream_end))
        except Exception as e:
            self.logger.error(f"Error streaming response for ID {request_id}: {e}")
            stream_error = {
                "type": "stream-error",
                "id": request_id,
                "error": str(e),
            }
            try:
                await websocket.send(json.dumps(stream_error))
            except Exception:
                pass
        finally:
            self.active_streams.pop(request_id, None)

    def run(self) -> None:
        try:
            asyncio.run(self.connect())
        except KeyboardInterrupt:
            self.logger.debug("Received exit signal")
