import os
import subprocess
import cv2
from ultralytics import YOLO

from exercises import EXERCISES

BEEP_FILE  = os.path.join(os.path.dirname(__file__), 'beep.mp3')
MODEL_PATH = os.path.join(os.path.dirname(__file__), 'yolo11x-pose.pt')

HINT = '[P]ush [U]pull [B]ench [C]url [S]itup [L]plank  [R]eset [Q]uit'


def beep():
    subprocess.Popen(['afplay', BEEP_FILE],
                     stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


model = YOLO(MODEL_PATH)

cap   = cv2.VideoCapture(0)
count = 0
stage = None
mode  = 'p'

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    results = model(frame, verbose=False)
    result  = results[0]

    # Draw YOLO skeleton on a copy of the frame (no bounding boxes or conf labels)
    image = result.plot(boxes=False, conf=False, labels=False)

    # Extract the first detected person's keypoints
    kps = conf = None
    if (result.keypoints is not None
            and result.keypoints.xyn is not None
            and len(result.keypoints.xyn) > 0):
        kps  = result.keypoints.xyn[0].cpu().numpy()   # (17, 2) normalised
        conf = result.keypoints.conf[0].cpu().numpy()   # (17,)

    if kps is not None:
        try:
            ex = EXERCISES[mode]
            new_stage, counted = ex['fn'](kps, conf, stage)
            if new_stage != stage:
                stage = new_stage
                beep()
            if counted:
                count += 1
        except Exception:
            pass

    # ── UI overlay ───────────────────────────────────────────────────────────
    h, w  = image.shape[:2]
    ex    = EXERCISES[mode]
    color = ex['color']

    cv2.rectangle(image, (0, 0), (w, 80), (15, 15, 15), -1)
    cv2.putText(image, ex['name'], (12, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.85, color, 2)
    cv2.putText(image, f'{ex["unit"]}: {count}', (12, 68),
                cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3)
    cv2.putText(image, (stage or '---').upper(), (w - 108, 50),
                cv2.FONT_HERSHEY_SIMPLEX, 0.9, (180, 180, 0), 2)

    cv2.rectangle(image, (0, h - 30), (w, h), (15, 15, 15), -1)
    cv2.putText(image, HINT, (8, h - 8),
                cv2.FONT_HERSHEY_SIMPLEX, 0.48, (160, 160, 160), 1)

    cv2.imshow('Exercise Counter', image)

    key = cv2.waitKey(1) & 0xFF
    if key == ord('q'):
        break
    elif key == ord('r'):
        count = 0
        stage = None
    elif key != 255 and chr(key) in EXERCISES:
        mode  = chr(key)
        stage = None
        count = 0

cap.release()
cv2.destroyAllWindows()