app.py

import asyncio
import json
import os
import shutil
import subprocess
import time
import logging
import concurrent.futures
import threading
import numpy as np
import uuid
from aiohttp import web
from aiortc import RTCPeerConnection, RTCSessionDescription, VideoStreamTrack, RTCIceServer, RTCConfiguration
from av import VideoFrame
import mss

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

# Xvfb buffer limits
MAX_WIDTH = 3840
MAX_HEIGHT = 2160

# Initial Resolution
DEFAULT_WIDTH = 1280
DEFAULT_HEIGHT = 720

# Cloudflare TURN Credentials
TURN_USER = "g08abe68c81a07f098bb5f0914549bb32440e5aad0b216c7fba2b61e76fd62c6"
TURN_PASS = "aed1a10dd10eba9401ad9d99e5c66036d8a970eab5ba8e6dc9845ab57c771a7d"

logging.basicConfig(level=logging.WARNING) 
logger = logging.getLogger("WebRTC-Brave")

# Increase worker threads for faster I/O
executor = concurrent.futures.ThreadPoolExecutor(max_workers=16)
thread_local_storage = threading.local()

config = {
    "width": DEFAULT_WIDTH,
    "height": DEFAULT_HEIGHT
}

# --- High Performance Input Manager ---
class InputManager:
    def __init__(self):
        self.process = None
        self.lock = threading.Lock()
        self.scroll_accum = 0
        # FIX: Do not start process immediately in __init__
        # It must wait for Xvfb to be ready.

    def start_process(self):
        # Only start if DISPLAY var is set
        if not os.environ.get("DISPLAY"):
            return

        try:
            self.process = subprocess.Popen(
                ['xdotool', '-'], 
                stdin=subprocess.PIPE, 
                encoding='utf-8', 
                bufsize=0 
            )
        except Exception as e:
            logger.error(f"Failed to start xdotool: {e}")

    def _send_raw(self, command):
        # Lazy initialization / Restart if dead
        if self.process is None or self.process.poll() is not None:
            self.start_process()
        
        if self.process:
            try:
                self.process.stdin.write(command + "\n")
                self.process.stdin.flush()
            except Exception:
                # If write fails, force restart next time
                try: self.process.kill()
                except: pass
                self.process = None

    def send(self, command):
        with self.lock:
            self._send_raw(command)

    def scroll(self, dy):
        with self.lock:
            self.scroll_accum += dy
            THRESHOLD = 40 
            while self.scroll_accum >= THRESHOLD:
                self._send_raw("click 5")
                self.scroll_accum -= THRESHOLD
            while self.scroll_accum <= -THRESHOLD:
                self._send_raw("click 4")
                self.scroll_accum += THRESHOLD

    def mouse_move(self, x, y): self.send(f"mousemove {x} {y}")
    def mouse_down(self, btn): self.send(f"mousedown {btn}")
    def mouse_up(self, btn): self.send(f"mouseup {btn}")
    def click(self, btn, repeat=1): self.send(f"click --repeat {repeat} {btn}")
    def key_down(self, key): self.send(f"keydown {key}")
    def key_up(self, key): self.send(f"keyup {key}")

input_manager = InputManager()

# --- System Management ---

def start_system():
    # 1. Setup Environment FIRST
    os.environ["DISPLAY"] = DISPLAY_NUM

    if not shutil.which("Xvfb"): raise FileNotFoundError("Xvfb missing")
    
    logger.warning(f"Starting Xvfb on {DISPLAY_NUM}...")
    subprocess.Popen([
        "Xvfb", DISPLAY_NUM, 
        "-screen", "0", f"{MAX_WIDTH}x{MAX_HEIGHT}x24",
        "-ac", "-noreset"
    ])
    
    # 2. Wait for Xvfb to initialize
    time.sleep(3)

    # 3. Now start xdotool (InputManager)
    input_manager.start_process()

    # 4. Initialize Resolution
    set_resolution(DEFAULT_WIDTH, DEFAULT_HEIGHT)

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

    # 6. Start Browser
    threading.Thread(target=keep_brave_alive, daemon=True).start()

def keep_brave_alive():
    brave_cmd = (
        "brave-browser "
        "--no-sandbox "
        "--start-maximized "
        "--user-data-dir=/home/user/brave-data "
        "--disable-infobars "
        "--disable-dev-shm-usage "
        "--disable-gpu "
        "--window-position=0,0 "
        f"--window-size={MAX_WIDTH},{MAX_HEIGHT}" 
    )
    while True:
        try:
            subprocess.run(brave_cmd, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            time.sleep(1)
        except: time.sleep(2)

def get_xrandr_output_name():
    try:
        out = subprocess.check_output(["xrandr"]).decode()
        for line in out.splitlines():
            if " connected" in line:
                return line.split()[0]
    except: pass
    return "screen"

def get_cvt_modeline(width, height, rate=60):
    H_BLANK = 160
    H_SYNC = 32
    H_FRONT_PORCH = 48
    V_FRONT_PORCH = 3
    V_SYNC = 5
    MIN_V_BLANK = 460
    
    frame_time_us = 1000000.0 / rate
    active_time_us = frame_time_us - MIN_V_BLANK
    if active_time_us <= 0: return None
    
    h_period_us = active_time_us / height
    v_blank_lines = int(MIN_V_BLANK / h_period_us) + 1
    v_total = height + v_blank_lines
    h_total = width + H_BLANK
    pclk = (h_total * v_total * rate) / 1000000.0
    
    h_sync_start = width + H_FRONT_PORCH
    h_sync_end = h_sync_start + H_SYNC
    v_sync_start = height + V_FRONT_PORCH
    v_sync_end = v_sync_start + V_SYNC
    
    return f'"{width}x{height}_60.00" {pclk:.2f} {width} {h_sync_start} {h_sync_end} {h_total} {height} {v_sync_start} {v_sync_end} {v_total} +hsync -vsync'

def set_resolution(w, h):
    try:
        if w % 2 != 0: w += 1
        if h % 2 != 0: h += 1

        output = get_xrandr_output_name()
        mode_name = f"WEB_{w}x{h}_{str(uuid.uuid4())[:4]}"
        modeline_str = get_cvt_modeline(w, h)
        if not modeline_str: return
            
        parts = modeline_str.split()
        mode_params = parts[1:]

        subprocess.run(["xrandr", "--newmode", mode_name] + mode_params, check=True)
        subprocess.run(["xrandr", "--addmode", output, mode_name], check=True)
        subprocess.run(["xrandr", "--output", output, "--mode", mode_name], check=True)
        
        config["width"] = w
        config["height"] = h
    except Exception as e:
        logger.error(f"Resolution setup failed: {e}")

# --- Video Capture ---
class VirtualScreenTrack(VideoStreamTrack):
    kind = "video"
    def __init__(self):
        super().__init__()
        self.last_frame_time = 0
        self.frame_count = 0

    def _capture(self):
        try:
            if not hasattr(thread_local_storage, "sct"):
                thread_local_storage.sct = mss.mss()
            
            monitor = {"top": 0, "left": 0, "width": config["width"], "height": config["height"]}
            sct_img = thread_local_storage.sct.grab(monitor)
            img = np.array(sct_img)
            return img[..., :3] 
        except: return None

    async def recv(self):
        FPS = 30
        FRAME_TIME = 1.0 / FPS
        
        pts, time_base = await self.next_timestamp()
        
        current_time = time.time()
        wait = FRAME_TIME - (current_time - self.last_frame_time)
        if wait > 0:
            await asyncio.sleep(wait)
        
        self.last_frame_time = time.time()

        frame = await asyncio.get_event_loop().run_in_executor(executor, self._capture)
        
        if frame is None:
            blank = np.zeros((config["height"], config["width"], 3), dtype=np.uint8)
            av_frame = VideoFrame.from_ndarray(blank, format="bgr24")
        else:
            av_frame = VideoFrame.from_ndarray(frame, format="bgr24")

        av_frame.pts = pts
        av_frame.time_base = time_base
        return av_frame

# --- Input Mapping ---
def map_key(key):
    if key == " ": return "space"
    k = key.lower()
    charmap = {
        "control": "ctrl", "shift": "shift", "alt": "alt", "meta": "super", "cmd": "super",
        "enter": "Return", "backspace": "BackSpace", "tab": "Tab", "escape": "Escape",
        "arrowup": "Up", "arrowdown": "Down", "arrowleft": "Left", "arrowright": "Right",
        "home": "Home", "end": "End", "pageup": "Page_Up", "pagedown": "Page_Down",
        "delete": "Delete", "insert": "Insert",
        "f1": "F1", "f2": "F2", "f3": "F3", "f4": "F4", "f5": "F5", "f6": "F6",
        "f7": "F7", "f8": "F8", "f9": "F9", "f10": "F10", "f11": "F11", "f12": "F12",
        "!": "exclam", "@": "at", "#": "numbersign", "$": "dollar", "%": "percent",
        "^": "asciicircum", "&": "ampersand", "*": "asterisk", "(": "parenleft",
        ")": "parenright", "-": "minus", "_": "underscore", "=": "equal", "+": "plus",
        "[": "bracketleft", "{": "braceleft", "]": "bracketright", "}": "braceright",
        ";": "semicolon", ":": "colon", "'": "apostrophe", "\"": "quotedbl",
        ",": "comma", "<": "less", ".": "period", ">": "greater", "/": "slash",
        "?": "question", "\\": "backslash", "|": "bar", "`": "grave", "~": "asciitilde",
        " ": "space"
    }
    return charmap.get(k, k)

def process_input(data):
    try:
        msg = json.loads(data)
        t = msg.get("type")
        
        current_w = config["width"]
        current_h = config["height"]

        if t == "resize":
            target_w = int(msg.get("width"))
            target_h = int(msg.get("height"))
            set_resolution(target_w, target_h)

        elif t == "mousemove":
            input_manager.mouse_move(int(msg["x"] * current_w), int(msg["y"] * current_h))
        elif t == "mousedown":
            input_manager.mouse_down({0:1, 1:2, 2:3}.get(msg.get("button"), 1))
        elif t == "mouseup":
            input_manager.mouse_up({0:1, 1:2, 2:3}.get(msg.get("button"), 1))
        elif t == "wheel":
            input_manager.scroll(msg.get("deltaY", 0))
        elif t == "keydown":
            k = map_key(msg.get("key"))
            if k: input_manager.key_down(k)
        elif t == "keyup":
            k = map_key(msg.get("key"))
            if k: input_manager.key_up(k)

    except Exception: pass 

# --- Routes ---
async def offer(request):
    try:
        params = await request.json()
        offer = RTCSessionDescription(sdp=params["sdp"], type=params["type"])
    except: return web.Response(status=400)

    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):
        channel.on("message", lambda m: asyncio.get_event_loop().run_in_executor(executor, process_input, m))

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

    sdp = "\r\n".join([l for l in pc.localDescription.sdp.splitlines() if "a=candidate" not in l or "typ relay" in l]) + "\r\n"
    return web.Response(content_type="application/json", text=json.dumps({"sdp": sdp, "type": pc.localDescription.type}), headers={"Access-Control-Allow-Origin": "*"})

async def index(r): return web.Response(text="helloworld")
async def options(r): return web.Response(headers={"Access-Control-Allow-Origin": "*", "Access-Control-Allow-Methods": "POST, OPTIONS", "Access-Control-Allow-Headers": "Content-Type"})

pcs = set()
async def on_shutdown(app): await asyncio.gather(*[pc.close() for pc in pcs])

if __name__ == "__main__":
    start_system()
    app = web.Application()
    app.on_shutdown.append(on_shutdown)
    app.router.add_get("/", index)
    app.router.add_post("/offer", offer)
    app.router.add_options("/offer", options)
    web.run_app(app, host=HOST, port=PORT)