import asyncio
import logging
import websockets

from config import TARGET_WS_URL, LOCAL_HOST, LOCAL_PORT
from selenium_utils import fetch_fresh_credentials

log = logging.getLogger(__name__)


async def handle_connection(local_client):
    log.info("Local client connected, fetching credentials...")

    result = await asyncio.get_running_loop().run_in_executor(
        None, fetch_fresh_credentials
    )
    if result is None:
        log.error("Could not fetch credentials")
        await local_client.close(1011, "Failed to bypass Cloudflare")
        return

    user_agent, cookies = result
    headers = {
        "User-Agent": user_agent,
        "Cookie": cookies,
    }

    try:
        log.info("Connecting to OmegleWeb...")
        async with websockets.connect(
            TARGET_WS_URL,
            additional_headers=headers,
        ) as target_server:
            log.info("OmegleWeb connected! Notifying client...")
            await local_client.send("SUCCESS")
            log.info("Starting relay")
            await relay_bidirectional(local_client, target_server)

    except websockets.exceptions.InvalidStatusCode as e:
        log.error(f"Omegle connection failed (Status {e.status_code})")
        await local_client.close(1011, f"Connection failed: {e.status_code}")

    except websockets.exceptions.ConnectionClosed:
        log.info("OmegleWeb disconnected")

    except Exception as e:
        log.error(f"Connection error: {type(e).__name__}: {e}")
        await local_client.close(1011, str(e))


async def relay_bidirectional(local_ws, target_ws):
    async def forward_local():
        try:
            async for message in local_ws:
                log.debug(f"[Local -> Omegle]: {message}")
                await target_ws.send(message)
        except websockets.exceptions.ConnectionClosed:
            log.info("Local client disconnected")
        except Exception as e:
            log.error(f"Forward local error: {e}")

    async def forward_remote():
        try:
            async for message in target_ws:
                log.debug(f"[Omegle -> Local]: {message}")
                await local_ws.send(message)
        except websockets.exceptions.ConnectionClosed:
            log.info("OmegleWeb disconnected")
        except Exception as e:
            log.error(f"Forward remote error: {e}")

    await asyncio.gather(
        forward_local(),
        forward_remote(),
    )


async def main():
    log.info(f"Starting WebSocket proxy on ws://{LOCAL_HOST}:{LOCAL_PORT}")
    await websockets.serve(handle_connection, LOCAL_HOST, LOCAL_PORT)
    await asyncio.Future()


if __name__ == "__main__":
    asyncio.run(main())