import asyncio
import json
import signal
import ssl
import base64
from contextlib import suppress
from pathlib import Path
import cv2
import websockets
from websockets.exceptions import ConnectionClosed

HOST = "0.0.0.0"
PORT = 8765

CERT_FILE = "server-cert.pem"
KEY_FILE = "server-key.pem"

CAMERA_INDEX = 0
FRAME_WIDTH = 1920
FRAME_HEIGHT = 1080

VIDEO_JPEG_QUALITY = 80
SNAPSHOT_JPEG_QUALITY = 95

FPS = 10
CAMERA_IDLE_TIMEOUT = 3

CAMERA_HUB = None


class CameraHub:
    def __init__(self) -> None:
        self.closed = False
        self.camera = None
        self.camera_running = False
        self.latest_frame = None
        self.latest_video_jpeg = None
        self.frame_id = 0
        self.frame_condition = asyncio.Condition()
        self.reader_task = None
        self.idle_close_task = None
        self.camera_lock = asyncio.Lock()
        self.camera_users = 0

    async def acquire_camera(self) -> None:
        """Megnyitja, illetve használatban tartja a kamerát."""
        async with self.camera_lock:
            if self.closed:
                raise RuntimeError("A kameraszerver már leállt.")

            self.camera_users += 1

            if self.idle_close_task:
                self.idle_close_task.cancel()
                self.idle_close_task = None

            if self.camera_running:
                return

            camera = await asyncio.to_thread(
                cv2.VideoCapture,
                CAMERA_INDEX,
            )
            camera.set(cv2.CAP_PROP_FRAME_WIDTH, FRAME_WIDTH)
            camera.set(cv2.CAP_PROP_FRAME_HEIGHT, FRAME_HEIGHT)

            if not camera.isOpened():
                camera.release()
                self.camera_users -= 1
                raise RuntimeError("A kamera nem nyitható meg.")

            self.camera = camera
            self.camera_running = True
            self.latest_frame = None
            self.latest_video_jpeg = None
            self.frame_id = 0
            self.reader_task = asyncio.create_task(self.read_frames())
            print("Kamera megnyitva.", flush=True)

    async def release_camera(self) -> None:
        """Elengedi egy kliens kamerahasználatát."""
        async with self.camera_lock:
            if self.camera_users > 0:
                self.camera_users -= 1

            if self.camera_users == 0 and self.camera_running:
                if self.idle_close_task:
                    self.idle_close_task.cancel()

                self.idle_close_task = asyncio.create_task(
                    self.close_camera_when_idle()
                )

    async def close_camera_when_idle(self) -> None:
        try:
            await asyncio.sleep(CAMERA_IDLE_TIMEOUT)

            async with self.camera_lock:
                if self.camera_users != 0 or not self.camera_running:
                    return

                await self.stop_camera()
                self.idle_close_task = None
        except asyncio.CancelledError:
            pass

    async def stop_camera(self) -> None:
        """Leállítja az olvasót és felszabadítja a kameraeszközt."""
        self.camera_running = False

        async with self.frame_condition:
            self.frame_condition.notify_all()

        if self.reader_task:
            self.reader_task.cancel()

            with suppress(asyncio.CancelledError):
                await self.reader_task

            self.reader_task = None

        if self.camera is not None:
            self.camera.release()
            self.camera = None

        self.latest_frame = None
        self.latest_video_jpeg = None
        print("Kamera bezárva (nincs aktív használó).", flush=True)

    async def read_frames(self) -> None:
        """
        Egyetlen háttérfeladatként olvassa a kamerát, és eltárolja
        a legutóbbi képkockát minden WebSocket kapcsolat számára.
        """
        frame_delay = 1 / FPS

        while not self.closed and self.camera_running:
            ok, frame = await asyncio.to_thread(self.camera.read)

            if not ok:
                await asyncio.sleep(0.05)
                continue

            video_jpeg = await self.encode_jpeg(
                frame,
                VIDEO_JPEG_QUALITY,
            )

            if video_jpeg is None:
                await asyncio.sleep(0.05)
                continue

            async with self.frame_condition:
                self.latest_frame = frame
                self.latest_video_jpeg = video_jpeg
                self.frame_id += 1
                self.frame_condition.notify_all()

            await asyncio.sleep(frame_delay)

    async def encode_jpeg(self, frame, quality: int):
        """
        JPEG-re kódolja a frame-et.
        """
        ok, jpeg = await asyncio.to_thread(
            cv2.imencode,
            ".jpg",
            frame,
            [cv2.IMWRITE_JPEG_QUALITY, quality],
        )

        if not ok:
            return None

        return jpeg.tobytes()

    async def wait_for_video_frame(self, last_frame_id: int):
        async with self.frame_condition:
            await self.frame_condition.wait_for(
                lambda: (
                    self.closed
                    or not self.camera_running
                    or self.frame_id != last_frame_id
                )
            )

            if self.closed or not self.camera_running:
                return None, None

            return self.frame_id, self.latest_video_jpeg

    async def send_frames(self, websocket, client_state) -> None:
        last_frame_id = 0

        while not self.closed:
            if not client_state["streaming"]:
                await asyncio.sleep(0.05)
                continue

            frame_id, jpeg_bytes = await self.wait_for_video_frame(
                last_frame_id,
            )

            if jpeg_bytes is None:
                break

            last_frame_id = frame_id

            try:
                await websocket.send(jpeg_bytes)
            except ConnectionClosed:
                break

    async def send_snapshot(
        self,
        websocket,
        request_id: str | None,
    ) -> None:
        """
        A legutóbbi frame-ből nagyobb minőségű fotót készít,
        és WebSocketen visszaküldi.
        """

        try:
            await self.acquire_camera()
        except RuntimeError as exc:
            await websocket.send(json.dumps({
                "type": "snapshot_error",
                "request_id": request_id,
                "message": str(exc),
            }))
            return

        try:
            # Előbb megpróbáljuk a legutóbbi frame-et használni.
            async with self.frame_condition:
                frame = (
                    None
                    if self.latest_frame is None
                    else self.latest_frame.copy()
                )

            # Ha még nincs frame, megvárjuk az elsőt.
            if frame is None:
                try:
                    await asyncio.wait_for(
                        self.wait_for_video_frame(0),
                        timeout=2,
                    )
                except asyncio.TimeoutError:
                    await websocket.send(
                        json.dumps({
                            "type": "snapshot_error",
                            "request_id": request_id,
                            "message": "Nem érkezett kép a kamerától.",
                        })
                    )
                    return

                async with self.frame_condition:
                    frame = (
                        None
                        if self.latest_frame is None
                        else self.latest_frame.copy()
                    )

            if frame is None:
                await websocket.send(
                    json.dumps({
                        "type": "snapshot_error",
                        "request_id": request_id,
                        "message": "Nem sikerült képet készíteni.",
                    })
                )
                return

            jpeg_bytes = await self.encode_jpeg(
                frame,
                SNAPSHOT_JPEG_QUALITY,
            )

            if jpeg_bytes is None:
                await websocket.send(
                    json.dumps({
                        "type": "snapshot_error",
                        "request_id": request_id,
                        "message": "Nem sikerült JPEG formátumba kódolni a képet.",
                    })
                )
                return

            jpeg_b64 = base64.b64encode(jpeg_bytes).decode("ascii")
            await websocket.send(json.dumps({
                "type": "snapshot",
                "image": "data:image/jpeg;base64," + jpeg_b64
            }))
        finally:
            await self.release_camera()

    async def receive_commands(self, websocket, client_state) -> None:
        async for message in websocket:
            if isinstance(message, bytes):
                continue

            try:
                data = json.loads(message)
            except json.JSONDecodeError:
                await websocket.send(
                    json.dumps({
                        "type": "error",
                        "message": "Érvénytelen JSON.",
                    })
                )
                continue

            command = data.get("cmd")

            if command == "start_video":
                client_state["streaming"] = True

                await websocket.send(
                    json.dumps({
                        "type": "status",
                        "streaming": True,
                    })
                )

            elif command == "stop_video":
                client_state["streaming"] = False

                await websocket.send(
                    json.dumps({
                        "type": "status",
                        "streaming": False,
                    })
                )

            elif command == "snapshot":
                await self.send_snapshot(
                    websocket,
                    data.get("request_id"),
                )

            elif command == "ping":
                await websocket.send(
                    json.dumps({
                        "type": "pong",
                    })
                )

            else:
                await websocket.send(
                    json.dumps({
                        "type": "error",
                        "message": f"Ismeretlen parancs: {command}",
                    })
                )

    async def close(self) -> None:
        self.closed = True

        async with self.camera_lock:
            if self.idle_close_task:
                self.idle_close_task.cancel()
                self.idle_close_task = None

            self.camera_users = 0

            if self.camera_running:
                await self.stop_camera()


async def websocket_handler(websocket) -> None:
    remote_address = websocket.remote_address
    print(f"WebSocket kapcsolat elfogadva: {remote_address}", flush=True)

    client_state = {
        "streaming": False,
    }
    camera_acquired = False

    try:
        await CAMERA_HUB.acquire_camera()
        camera_acquired = True

        await websocket.send(
            json.dumps({
                "type": "connected",
                "message": "Kapcsolódás sikeres",
            })
        )

        sender_task = asyncio.create_task(
            CAMERA_HUB.send_frames(websocket, client_state)
        )

        receiver_task = asyncio.create_task(
            CAMERA_HUB.receive_commands(websocket, client_state)
        )

        done, pending = await asyncio.wait(
            [sender_task, receiver_task],
            return_when=asyncio.FIRST_COMPLETED,
        )

        for task in pending:
            task.cancel()

        for task in pending:
            with suppress(asyncio.CancelledError):
                await task

        for task in done:
            exception = task.exception()

            if exception and not isinstance(
                exception,
                ConnectionClosed,
            ):
                raise exception

    except ConnectionClosed as exc:
        print(
            f"Kapcsolat bezárult: "
            f"code={exc.code}, reason={exc.reason}",
            flush=True,
        )

    except Exception as exc:
        print(f"Kapcsolati hiba: {exc!r}", flush=True)

    finally:
        client_state["streaming"] = False

        if camera_acquired:
            await CAMERA_HUB.release_camera()

        with suppress(ConnectionClosed):
            await websocket.close()

        print(f"Kapcsolat lezárva: {remote_address}", flush=True)


async def log_handshake_request(connection, request):
    print(
        "WebSocket upgrade kérés érkezett: "
        f"remote={connection.remote_address}, path={request.path}",
        flush=True,
    )

    return None


async def main() -> None:
    global CAMERA_HUB

    ssl_context = ssl.SSLContext(
        ssl.PROTOCOL_TLS_SERVER
    )

    ssl_context.load_cert_chain(
        certfile=CERT_FILE,
        keyfile=KEY_FILE,
    )

    stop_event = asyncio.Event()
    loop = asyncio.get_running_loop()

    for signal_name in (
        signal.SIGINT,
        signal.SIGTERM,
    ):
        with suppress(NotImplementedError):
            loop.add_signal_handler(
                signal_name,
                stop_event.set,
            )

    CAMERA_HUB = CameraHub()

    try:
        async with websockets.serve(
            websocket_handler,
            HOST,
            PORT,
            ssl=ssl_context,
            process_request=log_handshake_request,
            ping_interval=3,
            ping_timeout=3,
            close_timeout=1
        ):
            print(
                f"WSS szerver fut: "
                f"wss://0.0.0.0:{PORT}",
                flush=True,
            )

            await stop_event.wait()

    finally:
        await CAMERA_HUB.close()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass