"""Puente del Robot Virtual: hace que el robot parezca un ESP32-S3 en tu PC (ws://127.0.0.1:81).

MODO LOCAL (recomendado en casa): todo corre en tu PC, el video y los comandos no salen a internet.
    python puente.py --local
    1. Abre en el navegador de ESTA PC el enlace "modo local" de tu página del Robot Virtual.
    2. En la app Qt: IP 127.0.0.1, puerto 81 (como el ESP32).
    3. Video para OpenCV: http://127.0.0.1:8080/stream.mjpg  (o /frame.jpg)

MODO NUBE (celular como cámara): reenvía la app Qt al robot en la nube.
    python puente.py "wss://robot.bitgatelab.app/ws/robot?t=..."

Opciones: --puerto 81  --puerto-video 8080  --lan (aceptar conexiones de otros equipos de tu red,
por ejemplo la app Qt en tu celular Android; por omisión sólo esta PC).
Requiere: pip install "websockets>=13"
"""
import argparse
import asyncio
import json
import socket
import time

import websockets
from websockets.asyncio.server import serve

# Páginas que pueden conectar el mundo al puente. Las apps nativas (Qt, Python) no mandan Origin.
ORIGENES = {"https://robot.bitgatelab.app", "http://127.0.0.1:8787", "http://localhost:8787"}
MAX_TEXTO = 4096
MAX_CUADRO = 512 * 1024


def origen_permitido(connection, request):
    origin = request.headers.get("Origin")
    # Sin Origin = app nativa (Qt). Con Origin: el sitio del robot o algo de esta misma PC
    # (algunas librerías, p. ej. websocket-client de Python, mandan "http://127.0.0.1:81").
    local = origin is not None and origin.split("://")[-1].split(":")[0] in ("127.0.0.1", "localhost")
    if origin is not None and origin not in ORIGENES and not local:
        print(f"[!] conexión rechazada desde la página {origin}")
        return connection.respond(403, "Origen no permitido\n")
    return None


# ------------------------------------------------------------------ modo local

class Local:
    """Hace localmente lo mismo que el relay de la nube (src/session.js)."""

    def __init__(self):
        self.sims, self.ctls, self.cams = set(), set(), set()
        self.mjpeg = set()            # asyncio.Queue por cliente MJPEG
        self.ultimo = None            # último JPEG
        self.ultimo_t = 0.0
        self.ultimo_pedido = 0.0      # último GET /frame.jpg
        self.avisado = None

    # -- espectadores: el mundo sólo manda video si alguien lo ve
    def espectadores(self):
        return len(self.cams) + len(self.mjpeg) + (1 if time.time() - self.ultimo_pedido < 10 else 0)

    async def avisar(self, forzar=False):
        n = self.espectadores()
        if forzar or n != self.avisado:
            self.avisado = n
            await self.enviar(self.sims, json.dumps({"tipo": "_espectadores", "n": n}))

    @staticmethod
    async def enviar(destinos, msg):
        for ws in list(destinos):
            try:
                await ws.send(msg)
            except websockets.ConnectionClosed:
                pass

    # -- WebSocket: /sim = el mundo, /cam = video binario, cualquier otra ruta = la app (como el ESP32)
    async def ws(self, ws):
        ruta = ws.request.path.split("?")[0]
        if ruta.startswith("/sim"):
            await self.sesion_mundo(ws)
        elif ruta.startswith("/cam"):
            await self.sesion_cam(ws)
        else:
            await self.sesion_app(ws)

    async def sesion_mundo(self, ws):
        for viejo in list(self.sims):
            await viejo.close(4000, "Reemplazado por otra conexión del mundo")
        self.sims.add(ws)
        print("[+] mundo conectado")
        await self.enviar(self.ctls, json.dumps({"tipo": "robot_virtual", "conectado": True}))
        await self.avisar(forzar=True)
        try:
            async for msg in ws:
                if isinstance(msg, str):
                    if len(msg) <= MAX_TEXTO:
                        await self.enviar(self.ctls, msg)
                elif len(msg) <= MAX_CUADRO:
                    self.ultimo, self.ultimo_t = msg, time.time()
                    await self.enviar(self.cams, msg)
                    for q in list(self.mjpeg):
                        if q.full():
                            q.get_nowait()  # cliente lento: se descarta el cuadro viejo
                        q.put_nowait(msg)
                    await self.avisar()
        finally:
            self.sims.discard(ws)
            print("[-] mundo desconectado")
            if not self.sims:
                await self.enviar(self.ctls, json.dumps({"tipo": "robot_virtual", "conectado": False}))

    async def sesion_app(self, ws):
        self.ctls.add(ws)
        print(f"[+] app conectada desde {ws.remote_address[0]}")
        await ws.send(json.dumps({
            "tipo": "saludo", "mensaje": "Servidor ESP32-S3 listo", "puerto": 81,
            "nombre_host": "robot-virtual-local", "mac": "00:00:00:00:00:00",
            "virtual": True, "robot_conectado": bool(self.sims),
        }))
        try:
            async for msg in ws:
                if not isinstance(msg, str) or len(msg) > MAX_TEXTO:
                    continue
                if self.sims:
                    await self.enviar(self.sims, msg)
                else:
                    await ws.send(json.dumps({"tipo": "error", "detalle": "robot_virtual_desconectado"}))
        finally:
            self.ctls.discard(ws)
            print("[-] app desconectada")

    async def sesion_cam(self, ws):
        self.cams.add(ws)
        await self.avisar()
        try:
            await ws.wait_closed()
        finally:
            self.cams.discard(ws)
            await self.avisar()

    # -- HTTP de video: /stream.mjpg y /frame.jpg
    async def http(self, reader, writer):
        try:
            linea = (await reader.readline()).decode(errors="ignore")
            while (await reader.readline()) not in (b"\r\n", b"\n", b""):
                pass
            ruta = linea.split(" ")[1].split("?")[0] if " " in linea else "/"
            if ruta == "/stream.mjpg":
                await self.http_mjpeg(writer)
            elif ruta == "/frame.jpg":
                await self.http_frame(writer)
            else:
                writer.write(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
        except (ConnectionError, asyncio.IncompleteReadError):
            pass
        finally:
            try:
                writer.close()
            except Exception:
                pass

    async def http_frame(self, writer):
        dormido = self.espectadores() == 0
        self.ultimo_pedido = time.time()
        await self.avisar()
        if self.ultimo is None or (dormido and time.time() - self.ultimo_t > 2):
            writer.write(b"HTTP/1.1 503 Service Unavailable\r\nRetry-After: 1\r\nContent-Length: 0\r\n"
                         b"Access-Control-Allow-Origin: *\r\nConnection: close\r\n\r\n")
            return
        writer.write(b"HTTP/1.1 200 OK\r\nContent-Type: image/jpeg\r\nCache-Control: no-store\r\n"
                     b"Access-Control-Allow-Origin: *\r\n" + f"Content-Length: {len(self.ultimo)}\r\n".encode()
                     + b"Connection: close\r\n\r\n" + self.ultimo)
        await writer.drain()

    async def http_mjpeg(self, writer):
        writer.write(b"HTTP/1.1 200 OK\r\nContent-Type: multipart/x-mixed-replace; boundary=frame\r\n"
                     b"Cache-Control: no-store\r\nAccess-Control-Allow-Origin: *\r\nConnection: close\r\n\r\n")
        q = asyncio.Queue(maxsize=2)
        if self.ultimo:
            q.put_nowait(self.ultimo)
        self.mjpeg.add(q)
        await self.avisar()
        try:
            while True:
                jpg = await q.get()
                writer.write(b"--frame\r\nContent-Type: image/jpeg\r\n"
                             + f"Content-Length: {len(jpg)}\r\n\r\n".encode() + jpg + b"\r\n")
                await writer.drain()
        finally:
            self.mjpeg.discard(q)
            await self.avisar()


# ------------------------------------------------------------------ modo nube

async def tunel(cliente, url_nube: str):
    peer = cliente.remote_address
    print(f"[+] app conectada desde {peer[0]}")
    try:
        async with websockets.connect(url_nube, ping_interval=20, max_size=2**20) as nube:
            async def subir():
                async for msg in cliente:
                    await nube.send(msg)

            async def bajar():
                async for msg in nube:
                    if isinstance(msg, str):
                        await cliente.send(msg)

            tareas = [asyncio.create_task(subir()), asyncio.create_task(bajar())]
            _, pendientes = await asyncio.wait(tareas, return_when=asyncio.FIRST_COMPLETED)
            for t in pendientes:
                t.cancel()
    except websockets.InvalidStatus as e:
        code = e.response.status_code
        print(f"[!] la nube rechazó la conexión ({code}): " + ("el enlace venció o tu cuenta no está autorizada; genera uno nuevo en tu página" if code == 401 else "revisa la URL"))
        await cliente.close(4001, "token invalido")
    except OSError as e:
        print(f"[!] sin conexión con la nube: {e}")
        await cliente.close(1011, "sin nube")
    print(f"[-] app desconectada ({peer[0]})")


def ip_local() -> str:
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(("8.8.8.8", 80))
        return s.getsockname()[0]
    except OSError:
        return "127.0.0.1"
    finally:
        s.close()


async def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("url", nargs="?", help="modo nube: URL wss://.../ws/robot?t=... de tu página")
    ap.add_argument("--local", action="store_true", help="modo local: el mundo corre en el navegador de esta PC")
    ap.add_argument("--puerto", type=int, default=81)
    ap.add_argument("--puerto-video", type=int, default=8080)
    ap.add_argument("--lan", action="store_true", help="aceptar conexiones de otros equipos de la red")
    args = ap.parse_args()
    host = "0.0.0.0" if args.lan else "127.0.0.1"
    donde = f"127.0.0.1 (o {ip_local()} desde otro equipo)" if args.lan else "127.0.0.1"

    if args.local:
        loc = Local()
        async with serve(loc.ws, host, args.puerto, process_request=origen_permitido, max_size=MAX_CUADRO):
            video = await asyncio.start_server(loc.http, host, args.puerto_video)
            async with video:
                print("Puente LOCAL listo.")
                print(f"  1. Abre en el navegador de esta PC el enlace «modo local» de tu página.")
                print(f"  2. App Qt: IP {donde}, puerto {args.puerto}")
                print(f"  3. Video OpenCV: http://127.0.0.1:{args.puerto_video}/stream.mjpg")
                print("Ctrl+C para salir.")
                await asyncio.Future()
    else:
        if not args.url or "/ws/robot" not in args.url:
            raise SystemExit("Usa --local, o pasa la URL /ws/robot?t=... de tu página (modo nube).")
        async with serve(lambda ws: tunel(ws, args.url), host, args.puerto, process_request=origen_permitido):
            print(f"Puente a la NUBE listo. En la app Qt usa IP {donde} y puerto {args.puerto}.")
            print("Ctrl+C para salir.")
            await asyncio.Future()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass