from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse, StreamingResponse
import logging
import asyncio
import zmq.asyncio
import json
import datetime

from zmq import Socket
from control_backend.core.zmq_context import context
from control_backend.core.config import settings
from control_backend.schemas.ri_message import SpeechCommand, RIEndpoint




logger = logging.getLogger(__name__)

router = APIRouter()


@router.post("/command", status_code=202)
async def receive_command(command: SpeechCommand, request: Request):
    # Validate and retrieve data.
    SpeechCommand.model_validate(command)
    topic = b"command"
    pub_socket: Socket = request.app.state.internal_comm_socket
    pub_socket.send_multipart([topic, command.model_dump_json().encode()])

    return {"status": "Command received"}


@router.get("/ping_check")
async def ping(request: Request):
    pass


@router.get("/ping_stream")
async def ping_stream(request: Request):
    """Stream live updates whenever the device state changes."""
    async def event_stream():
        # Set up internal socket to receive ping updates
        logger.debug("Ping stream router event stream entered.")
        sub_socket = zmq.asyncio.Context().socket(zmq.SUB)
        sub_socket.connect(settings.zmq_settings.internal_comm_address)
        sub_socket.setsockopt(zmq.SUBSCRIBE, b"ping")
        connected = True

        ping_frequency = 1 # How many seconds between ping attempts

        # Even though its most likely the updates should alternate
        # So, True - False - True - False for connectivity.
        # Let's still check:)
        while True:
            logger.debug("Ping stream entered listening ")
            try:
                topic, body = await asyncio.wait_for(sub_socket.recv_multipart(), timeout=ping_frequency)
                logger.debug("got ping change in ping_stream router")
                connected = json.loads(body)
            except TimeoutError as e:
                await asyncio.sleep(0.1)
                
            # Stop if client disconnected
            if await request.is_disconnected():
                print("Client disconnected from SSE")
                break


            logger.debug(f"Yielded new connection event in robot ping router: {str(connected)}")
            yield f"data: {str(connected)}, time:{str(datetime.datetime.now().strftime("%H:%M:%S"))}\n\n"

            

    return StreamingResponse(event_stream(), media_type="text/event-stream")