import asyncio
import logging
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse, JSONResponse

from ros_node import get_node, start_ros, stop_ros

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

STATIC_DIR = Path(__file__).parent / 'static'


class ConnectionManager:
    def __init__(self) -> None:
        self.connections: list[WebSocket] = []
        self._loop: asyncio.AbstractEventLoop | None = None

    def set_loop(self, loop: asyncio.AbstractEventLoop) -> None:
        self._loop = loop

    async def connect(self, ws: WebSocket) -> None:
        await ws.accept()
        self.connections.append(ws)

    def disconnect(self, ws: WebSocket) -> None:
        if ws in self.connections:
            self.connections.remove(ws)

    def broadcast_from_thread(self, data: dict) -> None:
        if self._loop and self._loop.is_running():
            asyncio.run_coroutine_threadsafe(self._broadcast(data), self._loop)

    async def _broadcast(self, data: dict) -> None:
        dead = []
        for ws in self.connections:
            try:
                await ws.send_json(data)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.disconnect(ws)


manager = ConnectionManager()


@asynccontextmanager
async def lifespan(app: FastAPI):
    manager.set_loop(asyncio.get_event_loop())
    node = start_ros()
    node.on_joint_states_callbacks.append(manager.broadcast_from_thread)
    node.on_ultrasonic_callbacks.append(manager.broadcast_from_thread)
    yield
    stop_ros()


app = FastAPI(lifespan=lifespan)


@app.websocket('/ws')
async def websocket_endpoint(ws: WebSocket):
    await manager.connect(ws)
    node = get_node()
    try:
        while True:
            data = await ws.receive_json()
            msg_type = data.get('type')
            if msg_type == 'cmd_vel':
                node.publish_cmd_vel(data.get('linear_x', 0.0), data.get('angular_z', 0.0))
            elif msg_type == 'joint_command':
                node.publish_joint_command(data.get('pan', 0.0), data.get('tilt', 0.0))
    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.error('WebSocket error: %s', e)
    finally:
        manager.disconnect(ws)
        try:
            node.publish_cmd_vel(0.0, 0.0)
        except Exception:
            pass


@app.get('/{full_path:path}')
async def serve_spa(full_path: str):
    index = STATIC_DIR / 'index.html'
    if not index.exists():
        return JSONResponse(
            {'detail': 'Frontend not built. In dev mode open http://localhost:5173'},
            status_code=404,
        )
    candidate = (STATIC_DIR / full_path).resolve()
    static_root = STATIC_DIR.resolve()
    if str(candidate).startswith(str(static_root)) and full_path and candidate.is_file():
        return FileResponse(candidate)
    return FileResponse(index)


if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host='0.0.0.0', port=8080)