import asyncio
import json
import os
import shutil
import subprocess
import time
import logging
from aiohttp import web
from aiortc import RTCPeerConnection, RTCSessionDescription, RTCIceServer, RTCConfiguration

# --- Configuration ---
HOST = "0.0.0.0"
PORT = 7860
DISPLAY_NUM = ":99"
VNC_PORT = 5900

# Initial Resolution
DEFAULT_WIDTH = 1280
DEFAULT_HEIGHT = 720

# Cloudflare TURN Credentials (Optional but recommended for remote access)
TURN_USER = "g08abe68c81a07f098bb5f0914549bb32440e5aad0b216c7fba2b61e76fd62c6"
TURN_PASS = "aed1a10dd10eba9401ad9d99e5c66036d8a970eab5ba8e6dc9845ab57c771a7d"

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("VNC-RTC-Bridge")

# --- System Process Management ---

def start_system():
    """Initializes the virtual display environment and VNC server."""
    os.environ["DISPLAY"] = DISPLAY_NUM
    
    # 1. Start Xvfb (Virtual Framebuffer)
    logger.info(f"Starting Xvfb on {DISPLAY_NUM}...")
    subprocess.Popen([
        "Xvfb", DISPLAY_NUM, 
        "-screen", "0", f"{DEFAULT_WIDTH}x{DEFAULT_HEIGHT}x24",
        "-ac", "-noreset"
    ])
    time.sleep(2) # Wait for X to initialize

    # 2. Start x11vnc (The VNC Server)
    # -forever: keep running after client disconnects
    # -shared: allow multiple connections
    # -nopw: no password for this internal bridge
    # -rfbport: VNC standard port
    logger.info(f"Starting x11vnc on port {VNC_PORT}...")
    subprocess.Popen([
        "x11vnc", "-display", DISPLAY_NUM,
        "-rfbport", str(VNC_PORT),
        "-forever", "-shared", "-nopw",
        "-bg", "-quiet", "-xkb"
    ])

    # 3. Start Window Manager (Matchbox)
    if shutil.which("matchbox-window-manager"):
        logger.info("Starting Matchbox Window Manager...")
        subprocess.Popen("matchbox-window-manager -use_titlebar no", shell=True)

    # 4. Start the Application (Opera)
    logger.info("Starting Opera Browser...")
    opera_cmd = (
        "opera "
        "--no-sandbox "
        "--start-maximized "
        "--user-data-dir=/home/user/opera-data "
        "--disable-infobars "
        "--disable-dev-shm-usage "
        "--disable-gpu "
        "--no-first-run "
    )
    subprocess.Popen(opera_cmd, shell=True)

# --- WebRTC DataChannel to VNC TCP Bridge ---

async def bridge_vnc_to_datachannel(channel):
    """
    Connects to the local VNC TCP port and pipes data bidirectionally 
    to the WebRTC DataChannel.
    """
    retry_count = 5
    reader, writer = None, None

    # Try to connect to the local VNC server
    while retry_count > 0:
        try:
            reader, writer = await asyncio.open_connection('127.0.0.1', VNC_PORT)
            break
        except ConnectionRefusedError:
            retry_count -= 1
            logger.warning(f"VNC connection refused, retrying... ({retry_count} left)")
            await asyncio.sleep(1)
    
    if not writer:
        logger.error("Could not connect to local VNC server.")
        channel.close()
        return

    logger.info("Successfully bridged DataChannel to local VNC.")

    # Task: Local VNC Socket (TCP) -> WebRTC DataChannel
    async def vnc_to_webrtc():
        try:
            while True:
                data = await reader.read(16384) # Read binary chunks
                if not data:
                    break
                channel.send(data)
        except Exception as e:
            logger.error(f"VNC to WebRTC Bridge error: {e}")
        finally:
            logger.info("VNC Socket closed.")

    # Task: WebRTC DataChannel -> Local VNC Socket (TCP)
    @channel.on("message")
    def on_message(message):
        if isinstance(message, bytes):
            # Write raw RFB protocol bytes to the VNC server
            writer.write(message)
        elif isinstance(message, str):
            # Handle potential JSON metadata (like manual resize requests)
            try:
                data = json.loads(message)
                if data.get("type") == "resize":
                    w, h = data["width"], data["height"]
                    subprocess.run(["xrandr", "--fb", f"{w}x{h}"], check=False)
            except:
                pass

    try:
        await vnc_to_webrtc()
    finally:
        writer.close()
        await writer.wait_closed()

# --- HTTP / WebRTC Handlers ---

pcs = set()

async def offer(request):
    """Handles the WebRTC signaling offer."""
    try:
        params = await request.json()
        offer = RTCSessionDescription(sdp=params["sdp"], type=params["type"])
    except Exception as e:
        return web.Response(status=400, text=str(e))

    # Configure WebRTC with STUN/TURN
    pc = RTCPeerConnection(RTCConfiguration(iceServers=[
        RTCIceServer(urls=["turns:turn.cloudflare.com:443?transport=tcp", "turn:turn.cloudflare.com:3478?transport=udp"], 
                     username=TURN_USER, credential=TURN_PASS),
        RTCIceServer(urls=["stun:stun.l.google.com:19302"]) 
    ]))
    pcs.add(pc)

    @pc.on("connectionstatechange")
    async def on_state():
        if pc.connectionState in ["failed", "closed"]:
            await pc.close()
            pcs.discard(pc)

    @pc.on("datachannel")
    def on_dc(channel):
        # When noVNC creates a DataChannel, start the TCP bridge
        asyncio.create_task(bridge_vnc_to_datachannel(channel))

    await pc.setRemoteDescription(offer)
    answer = await pc.createAnswer()
    await pc.setLocalDescription(answer)

    return web.Response(
        content_type="application/json", 
        text=json.dumps({"sdp": pc.localDescription.sdp, "type": pc.localDescription.type}), 
        headers={"Access-Control-Allow-Origin": "*"}
    )

async def options(request):
    """CORS preflight handler."""
    return web.Response(headers={
        "Access-Control-Allow-Origin": "*",
        "Access-Control-Allow-Methods": "POST, OPTIONS",
        "Access-Control-Allow-Headers": "Content-Type"
    })

async def on_shutdown(app):
    """Clean up PeerConnections on server stop."""
    coros = [pc.close() for pc in pcs]
    await asyncio.gather(*coros)
    pcs.clear()

# --- Main Entry ---

if __name__ == "__main__":
    # Start the X11 and VNC backend
    start_system()

    # Start the WebRTC Signaling Server
    app = web.Application()
    app.on_shutdown.append(on_shutdown)
    
    app.router.add_get("/", lambda r: web.Response(text="WebRTC VNC Bridge is running."))
    app.router.add_post("/offer", offer)
    app.router.add_options("/offer", options)
    
    web.run_app(app, host=HOST, port=PORT)