import asyncio
import json
import os
import signal

import websockets
import gi

gi.require_version("Gst", "1.0")
gi.require_version("GstWebRTC", "1.0")
gi.require_version("GstSdp", "1.0")

from gi.repository import Gst, GstWebRTC, GstSdp


SIGNALING_URL = os.environ.get("SIGNALING_URL") or os.environ.get("CAMERA_SIGNALING_URL", "ws://127.0.0.1:3000")
ROOM_ID = os.environ.get("ROOM_ID", "jetson-1")
VIDEO_DEVICE = os.environ.get("VIDEO_DEVICE", "/dev/video0")
VIDEO_WIDTH = int(os.environ.get("VIDEO_WIDTH", "1920"))
VIDEO_HEIGHT = int(os.environ.get("VIDEO_HEIGHT", "1080"))
VIDEO_FRAMERATE = int(os.environ.get("VIDEO_FRAMERATE", "30"))
VIDEO_BITRATE = int(os.environ.get("VIDEO_BITRATE", "4000000"))


def default_pipeline_desc():
    return f"""
v4l2src device={VIDEO_DEVICE} !
video/x-raw,format=UYVY,width={VIDEO_WIDTH},height={VIDEO_HEIGHT},framerate={VIDEO_FRAMERATE}/1 !
nvvidconv !
video/x-raw(memory:NVMM),format=NV12 !
nvv4l2h264enc bitrate={VIDEO_BITRATE} insert-sps-pps=true iframeinterval=30 !
h264parse !
rtph264pay config-interval=1 pt=96 !
application/x-rtp,media=video,encoding-name=H264,payload=96 !
webrtcbin name=webrtc bundle-policy=max-bundle
"""


PIPELINE_DESC = os.environ.get("PIPELINE_DESC") or default_pipeline_desc()

ws = None
loop = None
pipeline = None
webrtc = None
stopping = False


def log(message):
    print(message, flush=True)


def send_json(data):
    if ws and loop and not stopping:
        asyncio.run_coroutine_threadsafe(ws.send(json.dumps(data)), loop)


def on_offer_created(promise, _, __):
    promise.wait()
    reply = promise.get_reply()
    offer = reply.get_value("offer")

    webrtc.emit("set-local-description", offer, Gst.Promise.new())

    send_json({
        "type": "offer",
        "offer": {
            "type": "offer",
            "sdp": offer.sdp.as_text(),
        },
    })
    log("offer sent")


def on_negotiation_needed(element):
    promise = Gst.Promise.new_with_change_func(on_offer_created, None, None)
    element.emit("create-offer", None, promise)


def request_offer(reason):
    if not webrtc or stopping:
        return
    log(f"creating offer: {reason}")
    on_negotiation_needed(webrtc)


def on_ice_candidate(element, mlineindex, candidate):
    send_json({
        "type": "candidate",
        "candidate": {
            "candidate": candidate,
            "sdpMLineIndex": mlineindex,
        },
    })


async def handle_message(message):
    data = json.loads(message)

    if data["type"] == "joined" and data.get("role") == "viewer":
        request_offer("viewer joined")

    elif data["type"] == "answer":
        sdp = data["answer"]["sdp"]
        _, sdpmsg = GstSdp.SDPMessage.new()
        GstSdp.sdp_message_parse_buffer(bytes(sdp.encode()), sdpmsg)

        answer = GstWebRTC.WebRTCSessionDescription.new(
            GstWebRTC.WebRTCSDPType.ANSWER,
            sdpmsg,
        )

        webrtc.emit("set-remote-description", answer, Gst.Promise.new())
        log("answer applied")

    elif data["type"] == "candidate":
        candidate = data["candidate"]
        webrtc.emit(
            "add-ice-candidate",
            candidate.get("sdpMLineIndex", 0),
            candidate["candidate"],
        )


async def main():
    global ws, loop, pipeline, webrtc

    loop = asyncio.get_running_loop()

    log(f"connecting signaling {SIGNALING_URL}, room={ROOM_ID}")
    log("pipeline:")
    log(PIPELINE_DESC)

    pipeline = Gst.parse_launch(PIPELINE_DESC)
    webrtc = pipeline.get_by_name("webrtc")
    if not webrtc:
        raise RuntimeError("webrtcbin named 'webrtc' was not found in pipeline.")

    webrtc.connect("on-negotiation-needed", on_negotiation_needed)
    webrtc.connect("on-ice-candidate", on_ice_candidate)

    async with websockets.connect(SIGNALING_URL) as websocket:
        ws = websocket

        await ws.send(json.dumps({
            "type": "join",
            "roomId": ROOM_ID,
            "role": "jetson",
        }))

        pipeline.set_state(Gst.State.PLAYING)
        log("pipeline started")
        await asyncio.sleep(0.5)
        request_offer("pipeline started")

        async for message in ws:
            await handle_message(message)


def stop_pipeline():
    global stopping
    stopping = True
    if pipeline:
        pipeline.set_state(Gst.State.NULL)
        log("pipeline stopped")


def handle_stop_signal(signum, frame):
    stop_pipeline()
    raise SystemExit(0)


if __name__ == "__main__":
    Gst.init(None)
    signal.signal(signal.SIGTERM, handle_stop_signal)
    signal.signal(signal.SIGINT, handle_stop_signal)

    try:
        asyncio.run(main())
    finally:
        stop_pipeline()