# app.py

import os
import argparse
import asyncio
import json
import ssl

from aiohttp import web, WSMsgType
import base64
from aiortc import RTCPeerConnection, RTCSessionDescription, VideoStreamTrack, AudioStreamTrack
from aiortc.contrib.media import MediaBlackhole, MediaRecorder
import cv2

from google import genai

# Set up Google GenAI client
client = genai.Client(http_options={"api_version": "v1alpha"})
CONFIG = {"generation_config": {"response_modalities": ["AUDIO"]}}
MODEL = "models/gemini-2.0-flash-exp"

pcs = set()

async def websocket_handler(request):
      ws = web.WebSocketResponse()
      await ws.prepare(request)

      # Store the websocket for communication
      ws_id = id(ws)
      request.app['websockets'][ws_id] = ws

      async for msg in ws:
          if msg.type == WSMsgType.TEXT:
              # Handle incoming messages if necessary
              pass
          elif msg.type == WSMsgType.ERROR:
              print(f'WebSocket connection closed with exception {ws.exception()}')

      # Clean up
      del request.app['websockets'][ws_id]
      return ws

async def index(request):
    content = open('index.html', 'r').read()
    return web.Response(content_type='text/html', text=content)

async def javascript(request):
    content = open('client.js', 'r').read()
    return web.Response(content_type='application/javascript', text=content)

async def offer(request):
    params = await request.json()
    offer = RTCSessionDescription(sdp=params['sdp'], type=params['type'])

    pc = RTCPeerConnection()
    pcs.add(pc)

    # Handle video track
    @pc.on('track')
    async def on_track(track):
        print('Track %s received' % track.kind)
        if track.kind == 'video':
            local_video = VideoFrameProcessorTrack(track, request.app)
            pc.addTrack(local_video)
        elif track.kind == 'audio':
            local_audio = AudioProcessorTrack(track, request.app)
            pc.addTrack(local_audio)
        pass

        @track.on('ended')
        async def on_ended():
            print('Track %s ended' % track.kind)

    await pc.setRemoteDescription(offer)

    # Send back an answer
    answer = await pc.createAnswer()
    await pc.setLocalDescription(answer)

    return web.Response(
        content_type='application/json',
        text=json.dumps(
            {'sdp': pc.localDescription.sdp, 'type': pc.localDescription.type}
        )
    )

async def on_shutdown(app):
    # Close all peer connections
    coros = [pc.close() for pc in pcs]
    await asyncio.gather(*coros)
    pcs.clear()

async def process_image_with_genai(image, app):
    # Convert image to bytes
    _, buffer = cv2.imencode('.jpg', image)
    image_bytes = buffer.tobytes()
    image_base64 = base64.b64encode(image_bytes).decode()

    frame_data = {"mime_type": "image/jpeg", "data": image_base64}

    # Use the Google GenAI client
    async with client.aio.live.connect(model=MODEL, config=CONFIG) as session:
        await session.send(frame_data)

        # Receive and process the response
        turn = session.receive()
        async for response in turn:
            if data := response.data:
                # Send data back to client via WebSocket
                message = {"type": "data", "data": base64.b64encode(data).decode()}
            if message_text := response.text:
                message = {"type": "text", "text": message_text}

            # Send message to all connected clients
            for ws in app['websockets'].values():
                await ws.send_json(message)

async def process_audio_with_genai(audio_data, app):
    # Prepare audio data
    audio_base64 = base64.b64encode(audio_data).decode()
    audio_frame = {"mime_type": "audio/pcm", "data": audio_base64}

    # Use the Google GenAI client
    async with client.aio.live.connect(model=MODEL, config=CONFIG) as session:
        await session.send(audio_frame)

        # Receive and process the response
        turn = session.receive()
        async for response in turn:
            if data := response.data:
                # Send data back to client via WebSocket
                message = {"type": "data", "data": base64.b64encode(data).decode()}
            if message_text := response.text:
                message = {"type": "text", "text": message_text}

            # Send message to all connected clients
            for ws in app['websockets'].values():
                await ws.send_json(message)

class VideoFrameProcessorTrack(VideoStreamTrack):
    """
    A video track that processes frames from another track.
    """
    def __init__(self, track, app):
        super().__init__()
        self.track = track
        self.app = app

    async def recv(self):
        frame = await self.track.recv()

        # Convert to OpenCV image
        img = frame.to_ndarray(format="bgr24")

        # Process the image with GenAI
        await process_image_with_genai(img, self.app)

        return frame

class AudioProcessorTrack(AudioStreamTrack):
     def __init__(self, track, app):
         super().__init__()
         self.track = track
         self.app = app

     async def recv(self):
         frame = await self.track.recv()
         audio_data = frame.to_bytes()

         # Process the audio frame with GenAI
         await process_audio_with_genai(audio_data, self.app)

         return frame  # Echo back if needed

def main():
    parser = argparse.ArgumentParser(description='WebRTC video stream')
    parser.add_argument('--host', default='0.0.0.0', help='Host to listen on (default: 0.0.0.0)')
    parser.add_argument('--port', type=int, default=8237, help='Port to listen on (default: 8237)')
    parser.add_argument('--cert', help='SSL certificate file (for HTTPS)')
    parser.add_argument('--key', help='SSL key file (for HTTPS)')
    args = parser.parse_args()

    app = web.Application()
    app['websockets'] = {}

    # Add the websocket route
    app.router.add_get('/ws', websocket_handler)
    app.on_shutdown.append(on_shutdown)
    app.router.add_get('/', index)
    app.router.add_get('/client.js', javascript)
    app.router.add_post('/offer', offer)

    # SSL context for HTTPS
    ssl_context = None
    if args.cert and args.key:
        ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
        ssl_context.load_cert_chain(args.cert, args.key)

    web.run_app(app, host=args.host, port=args.port, ssl_context=ssl_context)

if __name__ == '__main__':
    main()