# tracking_data.py
import asyncio
import json
from fastapi import WebSocket, WebSocketDisconnect
from aiokafka import AIOKafkaProducer
from datetime import datetime

producer: AIOKafkaProducer = None
connected_clients = set()

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()

# This endpoint receives tracking data and pushes it into Kafka
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

# This endpoint keeps clients connected for real-time map updates
async def websocket_realtime_updates(websocket: WebSocket):
    await websocket.accept()
    connected_clients.add(websocket)
    try:
        while True:
            await asyncio.sleep(0.1)
    except WebSocketDisconnect:
        connected_clients.remove(websocket)