from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from typing import List
from datetime import datetime
from bson import ObjectId
from app.db import database
from app.v1.dependencies.auth import verify_token  # If needed

router = APIRouter()
active_connections: List[WebSocket] = []

@router.websocket("/ws/tracking_data")
async def websocket_tracking_data(websocket: WebSocket, token: str = Query(...)):
    await websocket.accept()
    active_connections.append(websocket)

    # 🔧 Fix: Call the generator manually
    db_gen = database.get_mongo_db()
    db = next(db_gen)

    try:
        while True:
            data = await websocket.receive_json()

            telemetry = {
                "telemetry_id": str(ObjectId()),
                "device_id": data["device_id"],
                "timestamp": datetime.utcnow(),
                "location": data["location"],
                "speed": data.get("speed"),
                "heading": data.get("heading"),
                "payload": data.get("payload", {}),
                "created_date": datetime.utcnow(),
                "user_id": "eaafsaf"
            }

            db["tracking_data"].insert_one(telemetry)
            await websocket.send_json({"status": "ok", "message": "Telemetry stored."})

    except WebSocketDisconnect:
        active_connections.remove(websocket)
    finally:
        db_gen.close()  # ✅ Close generator manually