import argparse
import asyncio
import json
import struct
import time
import urllib.request

import numpy as np

FRAME_DIMENSION = 133
BLOCK_HEADER = struct.Struct("<III")


class MotionCollector:
    def __init__(self) -> None:
        self.started_at = time.perf_counter()
        self.first_block_seconds = None
        self.total_frames_received = 0
        self.blocks_received = 0
        self.last_block = None

    def elapsed(self) -> float:
        return time.perf_counter() - self.started_at

    def consume(self, message) -> bool:
        if isinstance(message, (bytes, bytearray)):
            request_index, frame_offset, frame_count = BLOCK_HEADER.unpack_from(message, 0)
            expected_length = BLOCK_HEADER.size + frame_count * FRAME_DIMENSION * 4
            if len(message) != expected_length:
                print("Bad block length", len(message), "expected", expected_length)
                return True

            frames = np.frombuffer(message, dtype="<f4", offset=BLOCK_HEADER.size)
            frames = frames.reshape(frame_count, FRAME_DIMENSION)

            if self.first_block_seconds is None:
                self.first_block_seconds = self.elapsed()

            self.total_frames_received += frame_count
            self.blocks_received += 1
            self.last_block = frames
            return False

        data = json.loads(message)
        message_type = data.get("type")

        if message_type == "hello":
            print("hello:", "ready =", data.get("ready"), "fps =", data.get("fps"), "transport =", data.get("transport"))
            return False

        if message_type == "sentence":
            print(
                "sentence", data["sentence_index"] + 1, "of", data["sentence_count"],
                "frames =", data["frame_count"],
                "generation_ms =", data["generation_ms"],
                "cache_hit =", data["cache_hit"],
                "text =", data["text"],
            )
            return False

        if message_type == "done":
            print("done: total_frames =", data["total_frames"], "server_elapsed_ms =", data["elapsed_ms"])
            return True

        if message_type == "error":
            print("error:", data)
            return True

        print("message:", data)
        return False

    def report(self) -> None:
        print("blocks received:", self.blocks_received)
        print("frames received:", self.total_frames_received)
        if self.first_block_seconds is not None:
            print("seconds until first motion block:", round(self.first_block_seconds, 3))
        print("seconds until finished:", round(self.elapsed(), 3))
        if self.last_block is not None:
            print("last block shape:", self.last_block.shape)
            print("last block value range:", float(self.last_block.min()), float(self.last_block.max()))


def build_generate_message(text: str) -> str:
    return json.dumps({"type": "generate", "id": "test-1", "text": text})


def http_request(url: str, token: str, payload=None) -> dict:
    headers = {"Content-Type": "application/json"}
    if token:
        headers["Authorization"] = "Bearer " + token

    request_body = None
    if payload is not None:
        request_body = json.dumps(payload).encode("utf-8")

    request = urllib.request.Request(url, data=request_body, headers=headers)
    with urllib.request.urlopen(request, timeout=60) as response:
        return json.loads(response.read().decode("utf-8"))


async def run_websocket_test(base_url: str, text: str, token: str) -> None:
    import websockets

    websocket_url = base_url.replace("http", "ws", 1) + "/ws"
    if token:
        websocket_url += "?token=" + token

    print("WEBSOCKET TEST", websocket_url)
    collector = MotionCollector()

    async with websockets.connect(websocket_url, max_size=None) as socket:
        await socket.send(build_generate_message(text))
        collector.started_at = time.perf_counter()

        while True:
            message = await asyncio.wait_for(socket.recv(), timeout=180)
            if collector.consume(message):
                break

    collector.report()


async def run_webrtc_test(base_url: str, text: str, token: str) -> None:
    from aiortc import RTCConfiguration, RTCIceServer, RTCPeerConnection, RTCSessionDescription

    print("WEBRTC TEST", base_url)

    ice_response = await asyncio.to_thread(http_request, base_url + "/ice", token)
    ice_servers = []
    for entry in ice_response["iceServers"]:
        ice_servers.append(
            RTCIceServer(
                urls=entry["urls"],
                username=entry.get("username"),
                credential=entry.get("credential"),
            )
        )

    peer = RTCPeerConnection(configuration=RTCConfiguration(iceServers=ice_servers))
    channel = peer.createDataChannel("motion", ordered=True)
    collector = MotionCollector()
    finished = asyncio.get_event_loop().create_future()

    @channel.on("open")
    def on_open() -> None:
        print("data channel open")
        collector.started_at = time.perf_counter()
        channel.send(build_generate_message(text))

    @channel.on("message")
    def on_message(message) -> None:
        if collector.consume(message) and not finished.done():
            finished.set_result(True)

    offer = await peer.createOffer()
    await peer.setLocalDescription(offer)

    answer = await asyncio.to_thread(
        http_request,
        base_url + "/webrtc/offer",
        token,
        {"sdp": peer.localDescription.sdp, "type": peer.localDescription.type},
    )

    await peer.setRemoteDescription(RTCSessionDescription(sdp=answer["sdp"], type=answer["type"]))

    try:
        await asyncio.wait_for(finished, timeout=180)
    finally:
        await peer.close()

    collector.report()


async def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base", default="http://127.0.0.1:7860")
    parser.add_argument("--text", default="Hello. How are you today? The doctor will help you now.")
    parser.add_argument("--mode", choices=["websocket", "webrtc", "both"], default="both")
    parser.add_argument("--token", default="")
    arguments = parser.parse_args()

    if arguments.mode in ("websocket", "both"):
        await run_websocket_test(arguments.base, arguments.text, arguments.token)

    if arguments.mode in ("webrtc", "both"):
        await run_webrtc_test(arguments.base, arguments.text, arguments.token)


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