import asyncio
import logging
import os
from uuid import uuid4

from fastapi import WebSocket, WebSocketDisconnect
from starlette.authentication import requires

import conf
import ttun_server
from proxy.queue import ProxyQueue
from ttun_server.types import (
    Config,
    Message,
    MessageType,
)

logger = logging.getLogger(__name__)
logger.setLevel("DEBUG")


async def assert_compatible_version(websocket: WebSocket, config: Config) -> None:
    client_version = config.get("version", "1.0.0")
    logger.debug("client_version %s", client_version)

    if "git" not in client_version and ttun_server.__version__ != "development":
        [client_major, *_] = [int(i) for i in client_version.split(".")[:3]]
        [server_major, *_] = [int(i) for i in ttun_server.__version__.split(".")]

        if client_major < server_major:
            await websocket.close(4000, "Your client is too old")

        if client_major > server_major:
            await websocket.close(4001, "Your client is too new")


@requires(["authenticated", "token"] if conf.ENABLE_AUTH else [])
async def tunnel(websocket: WebSocket) -> None:
    request_tasks: dict[str, asyncio.Task] = {}
    proxy_queues: dict[str, ProxyQueue] = {}

    await websocket.accept()
    config: Config = await websocket.receive_json()

    await assert_compatible_version(websocket, config)

    if "subdomains" not in config:
        config["subdomains"] = [config["subdomain"]]
    elif config["subdomains"] is None:
        config["subdomains"] = [None]

    for i, subdomain in enumerate(config["subdomains"]):
        if subdomain is None or await ProxyQueue.has_connection(subdomain):
            config["subdomains"][i] = uuid4().hex

    for subdomain in config["subdomains"]:
        proxy_queues[subdomain] = await ProxyQueue.create_for_identifier(subdomain)

    hostname = os.environ.get("TUNNEL_DOMAIN")
    protocol = "https" if os.environ.get("SECURE") else "http"

    urls = [
        f"{protocol}://{subdomain}.{hostname}" for subdomain in config["subdomains"]
    ]

    await websocket.send_json(
        {
            "url": urls[0],
            "urls": urls,
        }
    )

    async def handle_requests(subdomain: str) -> None:
        while request := await proxy_queues[subdomain].dequeue():
            asyncio.create_task(
                websocket.send_json(request), name=request["identifier"]
            )

    for subdomain in config["subdomains"]:
        request_tasks[subdomain] = asyncio.create_task(
            handle_requests(subdomain), name=subdomain
        )

    try:
        while True:
            data: Message = await websocket.receive_json()
            try:
                data["type"] = MessageType(data["type"]).value
                response_queue = await ProxyQueue.get_for_identifier(data["identifier"])
                await response_queue.enqueue(data)
            except AssertionError:
                pass
    except WebSocketDisconnect:
        for proxy_queue in proxy_queues.values():
            await proxy_queue.delete()
        for request_task in request_tasks.values():
            request_task.cancel()