import aiohttp
from aiohttp import web
import asyncio
from datetime import datetime

minirt_vnc_host = "localhost"
minirt_vnc_port = 7080

connections = {}
max_duration = 60


async def proxy_http(request):
    service_name = request.match_info.get('service')
    target = request.match_info.get('target')

    if service_name == 'minirt':
        target = f"{minirt_vnc_host}:{minirt_vnc_port}/{target}"
    elif service_name == 'minishell':
        target = ''

    print(f"proxying to host {target}")
    async with aiohttp.ClientSession() as session:
        async with session.request(
            request.method,
            f"http://{target}",
            headers=request.headers,
            data=await request.read()
        ) as resp:
            return web.Response(
                body=await resp.read(),
                status=resp.status,
                headers=resp.headers
            )

async def proxy_ws(request: web.Request):
    service_name = request.match_info.get('service')

    if service_name in connections:
        return web.Response(text='Service already in use', status=409)

    ws = web.WebSocketResponse()
    await ws.prepare(request)
    connections[service_name] = (ws, datetime.now())

    target = ''
    if service_name == 'minirt':
        target = f"{minirt_vnc_host}:{minirt_vnc_port}"
    elif service_name == 'minishell':
        target = ''

    try:
        async with aiohttp.ClientSession() as session:
            async with session.ws_connect(f"ws://{target}/websockify") as remote_ws:
                async def forward_to_remote():
                    async for msg in ws:
                        if msg.type == aiohttp.WSMsgType.BINARY:
                            await remote_ws.send_bytes(msg.data)

                async def forward_to_client():
                    async for msg in remote_ws:
                        if msg.type == aiohttp.WSMsgType.BINARY:
                            await ws.send_bytes(msg.data)

                async def check_timeout():
                    while True:
                        await asyncio.sleep(10)
                        elapsed = (datetime.now() - connections[service_name][1]).total_seconds()
                        if elapsed > max_duration:
                            await ws.close()
                            break

                await asyncio.gather(forward_to_remote(), forward_to_client(), check_timeout())
    finally:
        if service_name in connections:
            del connections[service_name]

    return ws