import asyncio
import json
from fastapi import WebSocket, WebSocketDisconnect
from aiokafka import AIOKafkaProducer
from datetime import datetime
from app.utils.connections import connected_clients

producer: AIOKafkaProducer = None


async def start_kafka_producer():
    global producer
    producer = AIOKafkaProducer(bootstrap_servers="localhost:9092")
    await producer.start()

async def stop_kafka_producer():
    if producer:
        await producer.stop()

async def websocket_tracking_data(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            message = await websocket.receive_text()
            data = json.loads(message)
            data["received_at"] = datetime.utcnow().isoformat()
            if producer:
                await producer.send("vehicle_tracking", json.dumps(data).encode("utf-8"))
            await websocket.send_text("Message queued successfully.")
    except WebSocketDisconnect:
        pass

async def websocket_realtime_updates(websocket: WebSocket):
    await websocket.accept()
    connected_clients.add(websocket)
    print("✅ Client connected. Total:", len(connected_clients))  # ✅ Moved here

    try:
        while True:
            await asyncio.sleep(1)
    except WebSocketDisconnect:
        connected_clients.remove(websocket)
        print("❌ Client disconnected. Remaining:", len(connected_clients))


