import cv2
import numpy as np
import gradio as gr
from gradio_webrtc import WebRTC

# Load face detector
face_cascade = cv2.CascadeClassifier(
    cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
)

def process_frame(frame: np.ndarray) -> np.ndarray:
    if frame is None:
        return frame

    frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
    gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)

    faces = face_cascade.detectMultiScale(
        gray, scaleFactor=1.1, minNeighbors=5, minSize=(60, 60)
    )

    for (x, y, w, h) in faces:
        mask_height = int(h * 0.4)
        mask_y = y + int(h * 0.2)

        # black mask
        cv2.rectangle(frame_bgr, (x, mask_y), (x + w, mask_y + mask_height), (0, 0, 0), -1)

        # eyes
        eye_y = mask_y + int(mask_height * 0.5)
        eye_radius = int(w * 0.1)

        cv2.circle(frame_bgr, (x + int(w * 0.35), eye_y), eye_radius, (255, 255, 255), -1)
        cv2.circle(frame_bgr, (x + int(w * 0.65), eye_y), eye_radius, (255, 255, 255), -1)

        # Z
        cv2.putText(
            frame_bgr, "Z",
            (x + w - 30, y + h - 20),
            cv2.FONT_HERSHEY_SIMPLEX,
            1.5,
            (0, 0, 255),
            3
        )

    return cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)


# UI
with gr.Blocks() as demo:
    gr.Markdown("# 🎭 Real-time Zorro Mask (WebRTC)")

    webrtc = WebRTC()

    webrtc.stream(
        fn=process_frame,
        inputs=webrtc,
        outputs=webrtc,
    )

demo.launch()