import asyncio
import websockets
import json
from fastapi import FastAPI
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

WEBSOCKET_URL = "wss://b917-160-179-178-105.ngrok-free.app/ws"
call_count = 0

async def websocket_client():
    global call_count
    while True:
        try:
            async with websockets.connect(WEBSOCKET_URL) as websocket:
                logger.info("WebSocket connection established")
                while True:
                    call_count += 1
                    payload = {"count": call_count}
                    await websocket.send(json.dumps(payload))
                    logger.info(f"Sent: {payload}")
                    response = await websocket.recv()
                    data = json.loads(response)
                    logger.info(f"Received: {data}")
                    await asyncio.sleep(1)  # Keep this delay as 1 second between messages
        except websockets.exceptions.ConnectionClosed:
            logger.error("WebSocket connection closed. Retrying...")
            await asyncio.sleep(0.1)  # 100ms delay before retrying
        except Exception as e:
            logger.error(f"Error: {e}")
            await asyncio.sleep(0.1)  # 100ms delay before retrying

@app.on_event("startup")
async def startup_event():
    asyncio.create_task(websocket_client())

@app.get("/")
async def root():
    return {"message": "WebSocket client is running", "calls_made": call_count}