import time
from av import VideoFrame
from aiortc import VideoStreamTrack
import numpy as np
from aiohttp import web
import argparse
import asyncio
import ssl
import os
from aiortc import RTCPeerConnection, RTCSessionDescription
import cv2
import json 
import time
from datetime import datetime

ROOT = "./"

class WebRTCVideoSink(VideoStreamTrack):
  def __init__(self, vr):
    """ vr - opencv is videocapture. """
    super().__init__()  # don't forget this!
    self.counter = 0
    self.vr = vr    

  async def recv(self):
#    time.sleep(1/40)
    ok, frame = self.vr.read()
    pts, time_base = await self.next_timestamp()
    if not ok:
      return None
    else:
      frame = VideoFrame.from_ndarray(frame, format="bgr24")
    frame.pts = pts
    frame.time_base = time_base
    self.counter += 1
    return frame

pcs_rgb = dict()

async def offer_video(request):
   params = await request.json()
   pc = RTCPeerConnection()
   pcs_rgb[request.remote] = pc
   reader = cv2.VideoCapture("demo.mp4")
   video_track = WebRTCVideoSink(reader)
   pc.addTrack(video_track)        
   offer = await pc.createOffer()
   print("Offer:", offer)
   await pc.setLocalDescription(offer)
   return web.Response(
            content_type="application/json",
            text=json.dumps(
                {"sdp": pc.localDescription.sdp, "type": pc.localDescription.type}
            ),
   )

async def answer_video(request):
   params = await request.json()
   answer = RTCSessionDescription(sdp=params["sdp"], type=params["type"])
   print("Answer:", answer)
   pc = pcs_rgb.get(request.remote, None)
   if not pc is None:
       await pc.setRemoteDescription(answer)
   return web.Response(text="Ok")

async def index(request):
   print(f"{datetime.now()}, {request.remote}") 
   content = open(os.path.join(ROOT, "index.html"), "r").read()
   return web.Response(content_type="text/html", text=content)

async def on_shutdown(app):
    # close peer connections
    coros = [pc.close() for remote, pc in pcs_rgb.items()]
    await asyncio.gather(*coros)
    pcs_rgb.clear()    

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="WebRTC audio / video / data-channels demo"
    )
    parser.add_argument("--cert-file", help="SSL certificate file (for HTTPS)", default="./certs/cert.pem")
    parser.add_argument("--key-file", help="SSL key file (for HTTPS)", default="./certs/key.pem")
    parser.add_argument(
        "--host", default="0.0.0.0", help="Host for HTTP server (default: 0.0.0.0)"
    )
    parser.add_argument(
        "--port", type=int, default=24000, help="Port for HTTP server (default: 8081)"
    )
    parser.add_argument("--verbose", "-v", action="count")
    args = parser.parse_args()
    
    # run event loop
    loop = asyncio.get_event_loop()
    
    app = web.Application()
    app.on_shutdown.append(on_shutdown)
    app.router.add_get("/", index)
    app.router.add_post("/offer_video", offer_video)
    app.router.add_post("/answer_video", answer_video)
    
    ssl_context = ssl.SSLContext()
    ssl_context.load_cert_chain(args.cert_file, args.key_file)
    
    web.run_app(
        app, access_log=None, host=args.host, port=args.port, ssl_context=ssl_context
    )