import asyncio
import websockets
import json


class WebSocketClient:
    def __init__(self, uri):
        self.uri = uri
        self.websocket = None
        self.is_connected = False

    async def connect(self):
        """Establishes the WebSocket connection."""
        if self.is_connected and self.websocket:
            print("Already connected.")
            return

        try:
            print(f"Connecting to {self.uri}...")
            self.websocket = await websockets.connect(self.uri)
            self.is_connected = True
            print("WebSocket connected.")
        except (ConnectionError, websockets.exceptions.ConnectionClosedOK) as e:
            print(f"Error connecting: {e}")
            self.is_connected = False
            self.websocket = None

    async def send_data(self, data):
        print(data)
        """Sends data over the open WebSocket connection."""
        if not self.is_connected or not self.websocket:
            print("WebSocket not connected. Attempting to reconnect...")
            await self.connect()
            if not self.is_connected:
                print("Failed to reconnect. Cannot send data.")
                return

        try:
            await self.websocket.send(json.dumps(data))
            print(f"Sent: {data}")
        except (ConnectionError, websockets.exceptions.ConnectionClosed) as e:
            print(f"Error sending data: {e}")
            self.is_connected = False
            self.websocket = None  # Reset connection on error
            await self.connect()  # Attempt to reconnect

    async def close(self):
        """Closes the WebSocket connection."""
        if self.websocket and self.is_connected:
            await self.websocket.close()
            self.is_connected = False
            self.websocket = None
            print("WebSocket closed.")