main.py

import asyncio
import json
import logging
import signal
from datetime import UTC, datetime, timedelta

import websockets
from aiortc import RTCPeerConnection, RTCSessionDescription
from requests import HTTPError

from intercomclient.camera_video_stream_track import CameraVideoStreamTrack
from intercomclient.config import Config
from intercomclient.device_authorization import (
    initiate_device_authorization,
    poll_for_token,
    refresh_tokens,
)
from intercomclient.token_store import TokenStatus, TokenStore

logging.basicConfig(level=logging.INFO)
LOG = logging.getLogger("pi-client")


class PiClient:
    def __init__(self, config: Config):
        self.config = config
        self.token_store = TokenStore(config)
        self.pc: RTCPeerConnection | None = None
        self.ws = None
        self.running = True

    # =========================
    # Token Management
    # =========================

    def check_token_status(self, tokens) -> TokenStatus:
        try:
            if tokens["access"]["token_value"] is None:
                return TokenStatus.MISSING
        except KeyError:
            return TokenStatus.MISSING

        if tokens["access"]["expiry_time"] < datetime.now(tz=UTC).timestamp():
            return TokenStatus.EXPIRED

        return TokenStatus.VALID

    async def ensure_valid_tokens(self):
        tokens = self.token_store.load_tokens()
        status = self.check_token_status(tokens)

        if status == TokenStatus.MISSING:
            LOG.info("No tokens found. Starting device authorization flow.")
            await asyncio.to_thread(self.device_authorization_flow)

        elif status == TokenStatus.EXPIRED:
            LOG.info("Access token expired. Attempting refresh.")
            await asyncio.to_thread(self.refresh_flow)

        else:
            expiry = datetime.fromtimestamp(tokens["access"]["expiry_time"], tz=UTC)
            LOG.info(
                "Access token valid, expires in %s",
                expiry - datetime.now(tz=UTC),
            )

    def device_authorization_flow(self):
        auth_response = initiate_device_authorization(self.config)
        LOG.info("User code: %s", auth_response["user_code"])

        token_response = poll_for_token(
            self.config,
            auth_response["device_code"],
            interval=auth_response["interval"],
        )

        print(token_response)

        expiry = (
            datetime.now(tz=UTC) + timedelta(seconds=token_response["expires_in"])
        ).timestamp()

        self.token_store.store_tokens(
            {
                "access_token": token_response["access_token"],
                "refresh_token": token_response["refresh_token"],
            },
            device_code=auth_response["device_code"],
            access_token_expiry=expiry,
        )

    def refresh_flow(self):
        tokens = self.token_store.load_tokens()

        try:
            refresh_response = refresh_tokens(
                self.config,
                tokens["refresh"]["token_value"],
            )

        except HTTPError as e:
            LOG.error("Token refresh failed: %s", e)
            print(refresh_response)

        expiry = (
            datetime.now(tz=UTC) + timedelta(seconds=refresh_response["expires_in"])
        ).timestamp()

        self.token_store.store_tokens(
            {
                "access_token": refresh_response["access_token"],
                "refresh_token": refresh_response["refresh_token"],
            },
            access_token_expiry=expiry,
        )

    # =========================
    # WebRTC + Signaling
    # =========================

    async def setup_peer_connection(self):
        self.pc = RTCPeerConnection()

        @self.pc.on("connectionstatechange")
        async def on_connectionstatechange():
            LOG.info("WebRTC state: %s", self.pc.connectionState)
            if self.pc.connectionState in ("failed", "closed"):
                await self.shutdown()

    async def signaling_loop(self):
        token_store = self.token_store.load_tokens()
        access_token = token_store["access"]["token_value"]
        device_code = token_store["device_code"]

        headers = {"Authorization": f"Bearer {access_token}"}

        websocket_url = f"{self.config.websocket_api_base_url}/{device_code}/"

        async with websockets.connect(
            websocket_url,
            additional_headers=headers,
        ) as ws:
            self.ws = ws
            LOG.info("Connected to signaling server")

            await self.setup_peer_connection()

            @self.pc.on("icecandidate")
            async def on_icecandidate(self, candidate):
                print(f"ICE candidate generated: {candidate}")
                if candidate is not None:
                    # Send the candidate to the signaling server
                    await self.ws.send_json(
                        json.dumps(
                            {
                                "type": "ice",
                                "candidate": {
                                    "candidate": candidate.component,
                                    "sdpMid": candidate.sdpMid,
                                    "sdpMLineIndex": candidate.sdpMLineIndex,
                                    "foundation": candidate.foundation,
                                    "priority": candidate.priority,
                                    "ip": candidate.address,
                                    "port": candidate.port,
                                    "protocol": candidate.protocol,
                                    "type": candidate.type,
                                },
                            }
                        )
                    )

            async for message in ws:
                data = json.loads(message)
                print(f"Received message: {data}")

                if data["type"] == "offer":
                    await self.pc.setRemoteDescription(
                        RTCSessionDescription(
                            sdp=data["sdp"],
                            type=data["type"],
                        )
                    )

                    camera_track = CameraVideoStreamTrack(self.config)
                    self.pc.addTrack(camera_track)
                    answer = await self.pc.createAnswer()

                    await self.pc.setLocalDescription(answer)
                    print(self.pc.localDescription.type)
                    await self.ws.send(
                        json.dumps(
                            {
                                "type": self.pc.localDescription.type,
                                "sdp": self.pc.localDescription.sdp,
                            }
                        )
                    )
                    print("Answer sent successfully")

                elif data["type"] == "ice":
                    await self.pc.addIceCandidate(data["candidate"])

    # =========================
    # Lifecycle
    # =========================

    async def run(self):
        while self.running:
            try:
                await self.ensure_valid_tokens()
                await self.signaling_loop()
            except Exception as e:
                LOG.exception("Client error: %s", e)
                await asyncio.sleep(5)

    async def shutdown(self):
        LOG.info("Shutting down client...")
        self.running = False

        if self.pc:
            await self.pc.close()

        if self.ws:
            await self.ws.close()


# =========================
# Entrypoint
# =========================


async def main():
    config = Config()
    client = PiClient(config)

    loop = asyncio.get_running_loop()

    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, lambda: asyncio.create_task(client.shutdown()))

    await client.run()


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