import asyncio
import websockets
import cv2
import numpy as np
from loguru import logger


async def handle_client(websocket):
    client = websocket.remote_address
    logger.info(f"Client connected: {client}")

    frame_count = 0

    try:
        async for message in websocket:
            frame_count += 1

            data = np.frombuffer(message, dtype=np.uint8)
            frame = cv2.imdecode(data, cv2.IMREAD_COLOR)

            if frame is None:
                logger.warning(f"Decode frame failed from client={client}")
                continue

            if frame_count % 100 == 0:
                logger.info(f"Received frames={frame_count}, client={client}, size={len(message)} bytes")

            cv2.imshow("Android Camera", frame)

            if cv2.waitKey(1) & 0xFF == 27:
                logger.info("ESC pressed, closing display")
                break

    except websockets.ConnectionClosed:
        logger.info(f"Client disconnected: {client}")

    except Exception as e:
        logger.exception(f"WebSocket error, client={client}, error={e}")

    finally:
        logger.info(f"Connection closed: client={client}, total_frames={frame_count}")
        cv2.destroyAllWindows()


async def main():
    host = "0.0.0.0"
    port = 8765

    logger.info(f"WebSocket server started: ws://{host}:{port}")

    async with websockets.serve(
        handle_client,
        host,
        port,
        max_size=10 * 1024 * 1024
    ):
        await asyncio.Future()


if __name__ == "__main__":
    asyncio.run(main())