import logging
from fastapi import APIRouter, WebSocket, WebSocketException
from ocpp_proto import chargepoint_manager
from ocpp_proto.chargepoint import ChargePoint
from util.websocket_wrapper import WebSocketWrapper

router = APIRouter()

@router.websocket("/{charging_station_id}")
async def websocket_endpoint(
    *,
    websocket: WebSocket,
    charging_station_id: str,
):
    """ For every new charging station that connects, create a ChargePoint
    instance and start listening for messages.
    """
    if (websocket.user.username != charging_station_id):
        raise WebSocketException(code=1008, reason="Username doesn't match chargepoint identifier")
    
    logging.info("Charging station '%s' connected", charging_station_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(charging_station_id, WebSocketWrapper(websocket))
    await chargepoint_manager.start(cp)