import logging
from fastapi import APIRouter, WebSocket, WebSocketException

from app.ocpp_proto import chargepoint_manager
from app.ocpp_proto.chargepoint import ChargePoint
from app.util.websocket_wrapper import WebSocketWrapper

router = APIRouter()

@router.websocket("/{chargepoint_identity}")
async def websocket_endpoint(
    *,
    websocket: WebSocket,
    chargepoint_identity: str,
):
    """ For every new charging station that connects, create a ChargePoint
    instance and start listening for messages.
    """
    if (websocket.user.identity != chargepoint_identity):
        raise WebSocketException(code=1008, reason="Username doesn't match chargepoint identifier")
    
    logging.info("Charging station '%s' (%s) connected", chargepoint_identity, websocket.user.id)

    # Check protocols
    try:
        requested_protocols = websocket.headers['sec-websocket-protocol']
        logging.info("Protocols advertised by charging station: %s", requested_protocols)
    except KeyError:
        logging.warning("Charging station hasn't advertised any subprotocol. "
                 "Closing Connection")
        return await websocket.close()

    if "ocpp2.0.1" in requested_protocols:
        logging.info("Matched supported protocol: ocpp2.0.1")
    else:
        logging.warning('Protocols mismatched | Expected subprotocols: %s,'
                        ' but client supports  %s | Closing connection',
                        "ocpp2.0.1",
                        requested_protocols)
        await websocket.accept()
        await websocket.close()
        return

    # Accept connection and begin communication
    await websocket.accept(subprotocol="ocpp2.0.1")
    cp = ChargePoint(chargepoint_identity, WebSocketWrapper(websocket))
    await chargepoint_manager.start(websocket.user.id, cp)