import uuid
from pathlib import Path

import cv2
import streamlit as st
from aiortc.contrib.media import MediaRecorder
from streamlit_webrtc import WebRtcMode, webrtc_streamer

from sample_utils.turn import get_ice_servers
from cvzone.HandTrackingModule import HandDetector
import numpy as np

RECORD_DIR = Path("./records")
RECORD_DIR.mkdir(exist_ok=True)

# Define virtual keyboard layout
keys = [
    ["Q", "W", "E", "R", "T", "Y", "U", "I", "O", "P"],
    ["A", "S", "D", "F", "G", "H", "J", "K", "L", ";"],
    ["Z", "X", "C", "V", "B", "N", "M", ",", ".", "/"],
]

# Class for keyboard buttons
class Button:
    def __init__(self, pos, text, size=(100, 100)):
        self.pos = pos
        self.size = size
        self.text = text

    def draw(self, img):
        x, y = self.pos
        w, h = self.size
        cv2.rectangle(img, (x, y), (x + w, y + h), (255, 0, 255), cv2.FILLED)
        cv2.putText(img, self.text, (x + 20, y + 70), cv2.FONT_HERSHEY_PLAIN, 3, (255, 255, 255), 3)

# Initialize buttons
button_list = []
for i, row in enumerate(keys):
    for j, key in enumerate(row):
        button_list.append(Button((j * 105 + 50, i * 120 + 50), key))


def video_frame_callback(frame):
    """
    Process video frames: Overlay a virtual keyboard on the video.
    """
    img = frame.to_ndarray(format="bgr24")

    # Draw virtual keyboard on the frame
    for button in button_list:
        button.draw(img)

    return av.VideoFrame.from_ndarray(img, format="bgr24")


def virtual_keyboard_app():
    """
    Streamlit application for displaying a virtual keyboard.
    """
    if "prefix" not in st.session_state:
        st.session_state["prefix"] = str(uuid.uuid4())
    prefix = st.session_state["prefix"]

    # Define input and output file paths
    in_file = RECORD_DIR / f"{prefix}_input.flv"
    out_file = RECORD_DIR / f"{prefix}_output.flv"

    # Factory for creating MediaRecorder objects
    def in_recorder_factory() -> MediaRecorder:
        return MediaRecorder(
            str(in_file), format="flv"
        )

    def out_recorder_factory() -> MediaRecorder:
        return MediaRecorder(str(out_file), format="flv")

    # Initialize WebRTC streamer
    webrtc_streamer(
        key="record",
        mode=WebRtcMode.SENDRECV,
        rtc_configuration={"iceServers": get_ice_servers()},
        media_stream_constraints={
            "video": True,
            "audio": True,
        },
        video_frame_callback=video_frame_callback,
        in_recorder_factory=in_recorder_factory,
        out_recorder_factory=out_recorder_factory,
    )

    # Display download buttons for recorded files
    if in_file.exists():
        with in_file.open("rb") as f:
            st.download_button(
                "Download the recorded video without video filter", f, "input.flv"
            )
    if out_file.exists():
        with out_file.open("rb") as f:
            st.download_button(
                "Download the recorded video with virtual keyboard", f, "output.flv"
            )


if __name__ == "__main__":
    virtual_keyboard_app()