from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Request, status
from fastapi.responses import JSONResponse
from uvicorn.logging import ColourizedFormatter

from rosbridge_library.internal.message_conversion import populate_instance, NonexistentFieldException, FieldTypeMismatchException
from toid_msgs.action import SimpleRotate, SimpleTranslateX, SimpleMoveCoords

import logging
import json

from toid_cli.services.services import ServerRunner

router = APIRouter()
log = logging.getLogger("ActionRoute")
log.addHandler(logging.StreamHandler())
log.handlers[0].setFormatter(ColourizedFormatter('%(levelprefix)s %(message)s'))
log.setLevel(logging.INFO)

@router.post("/start_tree")
async def start_tree(req: Request):
    tree = req.query_params.get("tree")
    ServerRunner.getRunner().execute_tree(tree)
    return JSONResponse("Works")

@router.post("/rotate")
async def rotate(req: Request):
    try:
        goal = await req.json()
    except json.decoder.JSONDecodeError:
        return JSONResponse("Bad Request", status_code=status.HTTP_400_BAD_REQUEST)
    log.info(goal)
    try:
        goal = populate_instance(goal, SimpleRotate.Goal())
        if not ServerRunner.getRunner().rotate_action(goal):
            return JSONResponse("Bad Request", status_code=status.HTTP_503_SERVICE_UNAVAILABLE)
    except (NonexistentFieldException, FieldTypeMismatchException) as e:
        return JSONResponse("Bad Request", status_code=status.HTTP_400_BAD_REQUEST)
    return JSONResponse("OK")

@router.post("/translate")
async def rotate(req: Request):
    try:
        goal = await req.json()
    except json.decoder.JSONDecodeError:
        return JSONResponse("Bad Request", status_code=status.HTTP_400_BAD_REQUEST)
    log.info(goal)
    try:
        goal = populate_instance(goal, SimpleMoveCoords.Goal())
        if not ServerRunner.getRunner().translate_coords(goal):
            return JSONResponse("Server busy", status_code=status.HTTP_503_SERVICE_UNAVAILABLE)
    except (NonexistentFieldException, FieldTypeMismatchException) as e:
        return JSONResponse("Bad Request", status_code=status.HTTP_400_BAD_REQUEST)
    return JSONResponse("OK")

@router.post("/translate_x")
async def rotate(req: Request):
    try:
        goal = await req.json()
    except json.decoder.JSONDecodeError:
        return JSONResponse("Bad Request", status_code=status.HTTP_400_BAD_REQUEST)
    log.info(goal)
    try:
        goal = populate_instance(goal, SimpleTranslateX.Goal())
        if not ServerRunner.getRunner().translate_x(goal):
            return JSONResponse("Server busy", status_code=status.HTTP_503_SERVICE_UNAVAILABLE)
    except (NonexistentFieldException, FieldTypeMismatchException) as e:
        return JSONResponse("Bad Request", status_code=status.HTTP_400_BAD_REQUEST)
    return JSONResponse("OK")

@router.post("/start_tf_pub")
async def startTFpub(req: Request):
    if not ServerRunner.getRunner().start_tf_pub():
        return JSONResponse("Server busy", status_code=status.HTTP_503_SERVICE_UNAVAILABLE)
    return JSONResponse("OK")


@router.websocket("/ws")
async def websocketHandler(ws: WebSocket):
    await ws.accept()
    try:
        ServerRunner.getRunner().add(ws)
        while True:
            data = await ws.receive_text()
            log.info(data)
            await ws.send_text("hello")
    except WebSocketDisconnect: 
        ServerRunner.getRunner().remove(ws)
        log.warning("User disconnected")
