import asyncio
import logging
import os
from uuid import uuid4

from fastapi import WebSocket, WebSocketDisconnect

import ttun_server
from ttun_server.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')


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', False) 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()