import json
from datetime import datetime
from aiokafka import AIOKafkaConsumer
from app.db import database
from app.v1.sockets.tracking_data import connected_clients  # ✅ Import WebSocket set

async def consume_tracking_data():
    consumer = AIOKafkaConsumer(
        "vehicle_tracking",
        bootstrap_servers='localhost:9092',
        group_id="tracking-group"
    )
    await consumer.start()

    mongo_gen = database.get_mongo_db()
    mongo = next(mongo_gen)
    collection = mongo["tracking_logs"]

    try:
        async for msg in consumer:
            data = json.loads(msg.value)
            data["received_at"] = datetime.utcnow()
            collection.insert_one(data)

            # ✅ Now inside the loop
            for ws in connected_clients.copy():
                try:
                    await ws.send_json(data)
                except Exception:
                    connected_clients.remove(ws)

            print("Broadcasting to", len(connected_clients), "clients")

    finally:
        await consumer.stop()
        try:
            next(mongo_gen)
        except StopIteration:
            pass