app.py

import os
import cv2
import numpy as np
import math
import base64
import json
import time
from pathlib import Path
from collections import deque
from typing import Dict
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from contextlib import asynccontextmanager
import torch
from ultralytics import YOLO

# Device detection - GPU if available, else CPU
DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using device: {DEVICE}")

# Model path - use yolov8n-pose which auto-downloads from ultralytics hub
MODEL_PATH = "yolov8n-pose.pt"

# --- Frame encoding settings ---
JPEG_QUALITY = 70
MAX_FRAME_DIM = 640


def _resize_frame(frame, max_dim=MAX_FRAME_DIM):
    """Resize frame if either dimension exceeds max_dim, preserving aspect ratio."""
    h, w = frame.shape[:2]
    if max(h, w) <= max_dim:
        return frame
    scale = max_dim / max(h, w)
    new_w, new_h = int(w * scale), int(h * scale)
    return cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_AREA)

# ============== PUSHUP MONITOR ==============
class PushupMonitor:
    def __init__(self, model_path=MODEL_PATH):
        self.model = YOLO(model_path)
        self.model.to(DEVICE)
        self.pushup_count = 0
        self.current_state = "UNKNOWN"
        self.prev_state = "UNKNOWN"
        self.angle_buffer = deque(maxlen=5)
        
        # Thresholds
        self.UP_ANGLE = 160
        self.DOWN_ANGLE = 90
    
    def process_frame(self, frame):
        results = self.model(frame, verbose=False, device=DEVICE)
        return results
    
    def calculate_angle(self, p1, p2, p3):
        v1 = np.array([p1[0] - p2[0], p1[1] - p2[1]])
        v2 = np.array([p3[0] - p2[0], p3[1] - p2[1]])
        cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2) + 1e-6)
        angle = np.degrees(np.arccos(np.clip(cos_angle, -1.0, 1.0)))
        return angle
    
    def get_smoothed_angle(self, angle):
        self.angle_buffer.append(angle)
        return np.mean(self.angle_buffer)
    
    def pushup_counter(self, results):
        if not results or len(results) == 0:
            return self.pushup_count, "UNKNOWN"
        
        result = results[0]
        if result.keypoints is None or len(result.keypoints) == 0:
            return self.pushup_count, "UNKNOWN"
        
        keypoints = result.keypoints.xy.cpu().numpy()
        confidence = result.keypoints.conf.cpu().numpy() if result.keypoints.conf is not None else None
        
        if keypoints.shape[0] == 0:
            return self.pushup_count, "UNKNOWN"
        
        kp = keypoints[0]
        conf = confidence[0] if confidence is not None else np.ones(17)
        
        # Use right side: shoulder(6), elbow(8), wrist(10)
        if conf[6] > 0.5 and conf[8] > 0.5 and conf[10] > 0.5:
            shoulder, elbow, wrist = kp[6], kp[8], kp[10]
        # Use left side: shoulder(5), elbow(7), wrist(9)
        elif conf[5] > 0.5 and conf[7] > 0.5 and conf[9] > 0.5:
            shoulder, elbow, wrist = kp[5], kp[7], kp[9]
        else:
            return self.pushup_count, "UNKNOWN"
        
        angle = self.calculate_angle(shoulder, elbow, wrist)
        smoothed = self.get_smoothed_angle(angle)
        
        if smoothed > self.UP_ANGLE:
            self.current_state = "UP"
        elif smoothed < self.DOWN_ANGLE:
            self.current_state = "DOWN"
        
        if self.prev_state == "DOWN" and self.current_state == "UP":
            self.pushup_count += 1
        
        self.prev_state = self.current_state
        return self.pushup_count, self.current_state
    
    def visualize(self, results):
        if not results or len(results) == 0:
            return np.zeros((480, 640, 3), dtype=np.uint8)
        
        frame = results[0].orig_img.copy()
        self.pushup_counter(results)
        
        # Draw skeleton
        if results[0].keypoints is not None:
            kp = results[0].keypoints.xy.cpu().numpy()
            if kp.shape[0] > 0:
                for point in kp[0]:
                    if point[0] > 0 and point[1] > 0:
                        cv2.circle(frame, (int(point[0]), int(point[1])), 4, (0, 255, 0), -1)
        
        return frame
    
    def reset_counter(self):
        self.pushup_count = 0
        self.current_state = "UNKNOWN"
        self.prev_state = "UNKNOWN"
        self.angle_buffer.clear()


# ============== SQUAT MONITOR ==============
class SquatMonitor:
    def __init__(self, model_path=MODEL_PATH):
        self.model = YOLO(model_path)
        self.model.to(DEVICE)
        self.squat_count = 0
        self.current_state = "UNKNOWN"
        self.prev_state = "UNKNOWN"
        self.angle_buffer = deque(maxlen=5)
        
        self.UP_ANGLE = 160
        self.DOWN_ANGLE = 90
    
    def process_frame(self, frame):
        return self.model(frame, verbose=False, device=DEVICE)
    
    def calculate_angle(self, p1, p2, p3):
        v1 = np.array([p1[0] - p2[0], p1[1] - p2[1]])
        v2 = np.array([p3[0] - p2[0], p3[1] - p2[1]])
        cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2) + 1e-6)
        return np.degrees(np.arccos(np.clip(cos_angle, -1.0, 1.0)))
    
    def get_smoothed_angle(self, angle):
        self.angle_buffer.append(angle)
        return np.mean(self.angle_buffer)
    
    def squat_counter(self, results):
        if not results or len(results) == 0:
            return self.squat_count, "UNKNOWN"
        
        result = results[0]
        if result.keypoints is None:
            return self.squat_count, "UNKNOWN"
        
        kp = result.keypoints.xy.cpu().numpy()
        conf = result.keypoints.conf.cpu().numpy() if result.keypoints.conf is not None else None
        
        if kp.shape[0] == 0:
            return self.squat_count, "UNKNOWN"
        
        kp, conf = kp[0], conf[0] if conf is not None else np.ones(17)
        
        # Right leg: hip(12), knee(14), ankle(16)
        if conf[12] > 0.5 and conf[14] > 0.5 and conf[16] > 0.5:
            hip, knee, ankle = kp[12], kp[14], kp[16]
        # Left leg: hip(11), knee(13), ankle(15)
        elif conf[11] > 0.5 and conf[13] > 0.5 and conf[15] > 0.5:
            hip, knee, ankle = kp[11], kp[13], kp[15]
        else:
            return self.squat_count, "UNKNOWN"
        
        angle = self.get_smoothed_angle(self.calculate_angle(hip, knee, ankle))
        
        if angle > self.UP_ANGLE:
            self.current_state = "UP"
        elif angle < self.DOWN_ANGLE:
            self.current_state = "DOWN"
        
        if self.prev_state == "DOWN" and self.current_state == "UP":
            self.squat_count += 1
        
        self.prev_state = self.current_state
        return self.squat_count, self.current_state
    
    def visualize(self, results):
        if not results or len(results) == 0:
            return np.zeros((480, 640, 3), dtype=np.uint8)
        frame = results[0].orig_img.copy()
        self.squat_counter(results)
        return frame
    
    def reset_counter(self):
        self.squat_count = 0
        self.current_state = "UNKNOWN"
        self.prev_state = "UNKNOWN"
        self.angle_buffer.clear()


# ============== CRUNCH MONITOR ==============
class CrunchMonitor:
    def __init__(self, model_path=MODEL_PATH):
        self.model = YOLO(model_path)
        self.model.to(DEVICE)
        self.crunch_count = 0
        self.current_state = "UNKNOWN"
        self.prev_state = "UNKNOWN"
        self.angle_buffer = deque(maxlen=5)
        
        self.UP_ANGLE = 60
        self.DOWN_ANGLE = 120
    
    def process_frame(self, frame):
        return self.model(frame, verbose=False, device=DEVICE)
    
    def calculate_angle(self, p1, p2, p3):
        v1 = np.array([p1[0] - p2[0], p1[1] - p2[1]])
        v2 = np.array([p3[0] - p2[0], p3[1] - p2[1]])
        cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2) + 1e-6)
        return np.degrees(np.arccos(np.clip(cos_angle, -1.0, 1.0)))
    
    def get_smoothed_angle(self, angle):
        self.angle_buffer.append(angle)
        return np.mean(self.angle_buffer)
    
    def crunch_counter(self, results):
        if not results or len(results) == 0:
            return self.crunch_count, "UNKNOWN"
        
        result = results[0]
        if result.keypoints is None:
            return self.crunch_count, "UNKNOWN"
        
        kp = result.keypoints.xy.cpu().numpy()
        conf = result.keypoints.conf.cpu().numpy() if result.keypoints.conf is not None else None
        
        if kp.shape[0] == 0:
            return self.crunch_count, "UNKNOWN"
        
        kp, conf = kp[0], conf[0] if conf is not None else np.ones(17)
        
        # shoulder(6), hip(12), knee(14)
        if conf[6] > 0.5 and conf[12] > 0.5 and conf[14] > 0.5:
            shoulder, hip, knee = kp[6], kp[12], kp[14]
        elif conf[5] > 0.5 and conf[11] > 0.5 and conf[13] > 0.5:
            shoulder, hip, knee = kp[5], kp[11], kp[13]
        else:
            return self.crunch_count, "UNKNOWN"
        
        angle = self.get_smoothed_angle(self.calculate_angle(shoulder, hip, knee))
        
        if angle < self.UP_ANGLE:
            self.current_state = "UP"
        elif angle > self.DOWN_ANGLE:
            self.current_state = "DOWN"
        
        if self.prev_state == "UP" and self.current_state == "DOWN":
            self.crunch_count += 1
        
        self.prev_state = self.current_state
        return self.crunch_count, self.current_state
    
    def visualize(self, results):
        if not results or len(results) == 0:
            return np.zeros((480, 640, 3), dtype=np.uint8)
        frame = results[0].orig_img.copy()
        self.crunch_counter(results)
        return frame
    
    def reset_counter(self):
        self.crunch_count = 0
        self.current_state = "UNKNOWN"
        self.prev_state = "UNKNOWN"
        self.angle_buffer.clear()


# ============== LEG RAISE MONITOR ==============
class LegRaiseMonitor:
    def __init__(self, model_path=MODEL_PATH):
        self.model = YOLO(model_path)
        self.model.to(DEVICE)
        self.rep_count = 0
        self.current_state = "UNKNOWN"
        self.prev_state = "UNKNOWN"
        self.height_buffer = deque(maxlen=5)
        
        self.UP_THRESHOLD = 0.6
        self.DOWN_THRESHOLD = 0.2
    
    def process_frame(self, frame):
        return self.model(frame, verbose=False, device=DEVICE)
    
    def leg_raise_counter(self, results):
        if not results or len(results) == 0:
            return self.rep_count, "UNKNOWN"
        
        result = results[0]
        if result.keypoints is None:
            return self.rep_count, "UNKNOWN"
        
        kp = result.keypoints.xy.cpu().numpy()
        conf = result.keypoints.conf.cpu().numpy() if result.keypoints.conf is not None else None
        
        if kp.shape[0] == 0:
            return self.rep_count, "UNKNOWN"
        
        kp, conf = kp[0], conf[0] if conf is not None else np.ones(17)
        
        # ankle(16), hip(12)
        if conf[16] > 0.5 and conf[12] > 0.5:
            ankle, hip = kp[16], kp[12]
        elif conf[15] > 0.5 and conf[11] > 0.5:
            ankle, hip = kp[15], kp[11]
        else:
            return self.rep_count, "UNKNOWN"
        
        # Calculate height ratio (higher ankle = smaller y = UP)
        height_ratio = (hip[1] - ankle[1]) / (hip[1] + 1e-6)
        self.height_buffer.append(height_ratio)
        smoothed = np.mean(self.height_buffer)
        
        if smoothed > self.UP_THRESHOLD:
            self.current_state = "UP"
        elif smoothed < self.DOWN_THRESHOLD:
            self.current_state = "DOWN"
        
        if self.prev_state == "UP" and self.current_state == "DOWN":
            self.rep_count += 1
        
        self.prev_state = self.current_state
        return self.rep_count, self.current_state
    
    def visualize(self, results):
        if not results or len(results) == 0:
            return np.zeros((480, 640, 3), dtype=np.uint8)
        frame = results[0].orig_img.copy()
        self.leg_raise_counter(results)
        return frame
    
    def reset_counter(self):
        self.rep_count = 0
        self.current_state = "UNKNOWN"
        self.prev_state = "UNKNOWN"
        self.height_buffer.clear()


# ============== PLANK MONITOR ==============
class PlankMonitor:
    def __init__(self, model_path=MODEL_PATH):
        self.model = YOLO(model_path)
        self.model.to(DEVICE)
        self.current_state = "REST"
        self.plank_start_time = None
        self.current_time = 0
        self.best_time = 0
        self.total_time = 0
        self.angle_buffer = deque(maxlen=5)
        
        self.MIN_PLANK_ANGLE = 150
        self.MAX_PLANK_ANGLE = 180
    
    def process_frame(self, frame):
        return self.model(frame, verbose=False, device=DEVICE)
    
    def calculate_angle(self, p1, p2, p3):
        v1 = np.array([p1[0] - p2[0], p1[1] - p2[1]])
        v2 = np.array([p3[0] - p2[0], p3[1] - p2[1]])
        cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2) + 1e-6)
        return np.degrees(np.arccos(np.clip(cos_angle, -1.0, 1.0)))
    
    def update_plank_state(self, results):
        if not results or len(results) == 0:
            self._handle_no_detection()
            return
        
        result = results[0]
        if result.keypoints is None:
            self._handle_no_detection()
            return
        
        kp = result.keypoints.xy.cpu().numpy()
        conf = result.keypoints.conf.cpu().numpy() if result.keypoints.conf is not None else None
        
        if kp.shape[0] == 0:
            self._handle_no_detection()
            return
        
        kp, conf = kp[0], conf[0] if conf is not None else np.ones(17)
        
        # shoulder(6), hip(12), ankle(16)
        if conf[6] > 0.5 and conf[12] > 0.5 and conf[16] > 0.5:
            shoulder, hip, ankle = kp[6], kp[12], kp[16]
        elif conf[5] > 0.5 and conf[11] > 0.5 and conf[15] > 0.5:
            shoulder, hip, ankle = kp[5], kp[11], kp[15]
        else:
            self._handle_no_detection()
            return
        
        angle = self.calculate_angle(shoulder, hip, ankle)
        self.angle_buffer.append(angle)
        smoothed = np.mean(self.angle_buffer)
        
        is_plank = self.MIN_PLANK_ANGLE <= smoothed <= self.MAX_PLANK_ANGLE
        
        if is_plank:
            if self.current_state != "PLANK":
                self.plank_start_time = time.time()
                self.current_state = "PLANK"
            else:
                self.current_time = time.time() - self.plank_start_time
                if self.current_time > self.best_time:
                    self.best_time = self.current_time
        else:
            if self.current_state == "PLANK":
                self.total_time += self.current_time
                self.current_time = 0
            self.current_state = "REST"
    
    def _handle_no_detection(self):
        if self.current_state == "PLANK":
            self.total_time += self.current_time
            self.current_time = 0
        self.current_state = "UNKNOWN"
    
    def visualize(self, results):
        if not results or len(results) == 0:
            return np.zeros((480, 640, 3), dtype=np.uint8)
        frame = results[0].orig_img.copy()
        self.update_plank_state(results)
        return frame
    
    def get_current_stats(self):
        return {
            "state": self.current_state,
            "current_time": round(self.current_time, 1),
            "best_time": round(self.best_time, 1),
            "total_time": round(self.total_time + self.current_time, 1)
        }
    
    def reset_counter(self):
        self.current_state = "REST"
        self.plank_start_time = None
        self.current_time = 0
        self.angle_buffer.clear()


# ============== BADDHA KONASANA (BUTTERFLY) MONITOR ==============
class BaddhaKonasanaMonitor:
    def __init__(self, model_path=MODEL_PATH):
        self.model = YOLO(model_path)
        self.model.to(DEVICE)

        self.L_SHOULDER, self.R_SHOULDER = 5, 6
        self.L_HIP, self.R_HIP = 11, 12
        self.L_KNEE, self.R_KNEE = 13, 14
        self.L_ANKLE, self.R_ANKLE = 15, 16

        self.in_pose = False
        self.pose_start_time = None
        self.current_hold_time = 0.0
        self.best_hold_time = 0.0
        self.total_hold_time = 0.0
        self.current_knee_spread = 0.0
        self.best_knee_spread = 0.0
        self.current_flexibility_score = 0.0

        self.MIN_KNEE_SPREAD = 60
        self.GOOD_KNEE_SPREAD = 120
        self.EXCELLENT_KNEE_SPREAD = 160
        self.MAX_FEET_DISTANCE = 100

        self.GRACE_SECONDS = 1.0
        self.last_valid_time = None
        self.spread_buffer = deque(maxlen=5)
        self.current_state = "READY"

    def process_frame(self, frame):
        return self.model(frame, verbose=False, device=DEVICE)

    def calculate_angle(self, p1, p2, p3):
        v1 = np.array([p1[0] - p2[0], p1[1] - p2[1]])
        v2 = np.array([p3[0] - p2[0], p3[1] - p2[1]])
        dot = np.dot(v1, v2)
        m1, m2 = np.linalg.norm(v1), np.linalg.norm(v2)
        if m1 * m2 == 0:
            return 180
        return math.degrees(math.acos(np.clip(dot / (m1 * m2), -1.0, 1.0)))

    def get_midpoint(self, p1, p2):
        return np.array([(p1[0] + p2[0]) / 2, (p1[1] + p2[1]) / 2])

    def update_pose_state(self, results):
        now = time.time()
        if not results or len(results) == 0:
            return self._handle_no_detection(now)
        result = results[0]
        if result.keypoints is None or len(result.keypoints) == 0:
            return self._handle_no_detection(now)

        kp = result.keypoints.xy.cpu().numpy()
        conf = result.keypoints.conf.cpu().numpy() if result.keypoints.conf is not None else None
        if kp.shape[0] == 0:
            return self._handle_no_detection(now)
        kp, conf = kp[0], conf[0] if conf is not None else np.ones(17)

        required = [self.L_HIP, self.R_HIP, self.L_KNEE, self.R_KNEE,
                    self.L_ANKLE, self.R_ANKLE, self.L_SHOULDER, self.R_SHOULDER]
        if not all(conf[i] > 0.4 for i in required):
            return self._handle_no_detection(now)

        mid_hip = self.get_midpoint(kp[self.L_HIP], kp[self.R_HIP])
        mid_shoulder = self.get_midpoint(kp[self.L_SHOULDER], kp[self.R_SHOULDER])
        knee_spread = self.calculate_angle(kp[self.L_KNEE], mid_hip, kp[self.R_KNEE])
        feet_dist = np.linalg.norm(kp[self.L_ANKLE] - kp[self.R_ANKLE])
        is_seated = mid_hip[1] > kp[self.L_SHOULDER][1]

        if knee_spread >= self.EXCELLENT_KNEE_SPREAD:
            flex = 100
        elif knee_spread >= self.GOOD_KNEE_SPREAD:
            flex = 70 + (knee_spread - self.GOOD_KNEE_SPREAD) / (self.EXCELLENT_KNEE_SPREAD - self.GOOD_KNEE_SPREAD) * 30
        elif knee_spread >= self.MIN_KNEE_SPREAD:
            flex = 30 + (knee_spread - self.MIN_KNEE_SPREAD) / (self.GOOD_KNEE_SPREAD - self.MIN_KNEE_SPREAD) * 40
        else:
            flex = (knee_spread / self.MIN_KNEE_SPREAD) * 30

        self.spread_buffer.append(knee_spread)
        smoothed = np.median(self.spread_buffer)
        self.current_knee_spread = smoothed
        self.current_flexibility_score = flex

        valid = knee_spread >= self.MIN_KNEE_SPREAD and is_seated and feet_dist < self.MAX_FEET_DISTANCE * 2
        if valid:
            if not self.in_pose:
                self.in_pose = True
                self.pose_start_time = now
                self.current_state = "HOLDING"
            self.last_valid_time = now
            self.best_knee_spread = max(self.best_knee_spread, smoothed)
            self.current_hold_time = now - self.pose_start_time
            self.best_hold_time = max(self.best_hold_time, self.current_hold_time)
        else:
            self._check_grace_period(now)

    def _handle_no_detection(self, now):
        self._check_grace_period(now)

    def _check_grace_period(self, now):
        if self.in_pose and self.last_valid_time is not None:
            if now - self.last_valid_time > self.GRACE_SECONDS:
                self.total_hold_time += self.current_hold_time
                self.in_pose = False
                self.pose_start_time = None
                self.current_hold_time = 0.0
                self.current_state = "READY"

    def visualize(self, results):
        if not results or len(results) == 0:
            return np.zeros((480, 640, 3), dtype=np.uint8)
        frame = results[0].orig_img.copy()
        self.update_pose_state(results)
        return frame

    def get_current_stats(self):
        return {
            "current_time": round(self.current_hold_time, 1),
            "best_time": round(self.best_hold_time, 1),
            "total_time": round(self.total_hold_time + self.current_hold_time, 1),
            "state": self.current_state,
        }

    def reset(self):
        self.in_pose = False
        self.pose_start_time = None
        self.current_hold_time = 0.0
        self.best_hold_time = 0.0
        self.total_hold_time = 0.0
        self.current_knee_spread = 0.0
        self.best_knee_spread = 0.0
        self.current_flexibility_score = 0.0
        self.current_state = "READY"
        self.spread_buffer.clear()


# ============== USTRASANA (CAMEL POSE) MONITOR ==============
class UstrasanaMonitor:
    def __init__(self, model_path=MODEL_PATH):
        self.model = YOLO(model_path)
        self.model.to(DEVICE)

        self.L_SHOULDER, self.R_SHOULDER = 5, 6
        self.L_HIP, self.R_HIP = 11, 12
        self.L_KNEE, self.R_KNEE = 13, 14

        self.in_pose = False
        self.pose_start_time = None
        self.current_hold_time = 0.0
        self.best_hold_time = 0.0
        self.total_hold_time = 0.0
        self.current_backbend_depth = 0.0
        self.best_backbend_depth = 0.0

        self.MIN_BACKBEND_ANGLE = 15
        self.GOOD_BACKBEND_ANGLE = 35
        self.EXCELLENT_BACKBEND_ANGLE = 55
        self.MIN_KNEE_HIP_DIFF = 50

        self.GRACE_SECONDS = 1.0
        self.last_valid_time = None
        self.depth_buffer = deque(maxlen=5)
        self.current_state = "READY"

    def process_frame(self, frame):
        return self.model(frame, verbose=False, device=DEVICE)

    def calculate_angle(self, p1, p2, p3):
        v1 = np.array([p1[0] - p2[0], p1[1] - p2[1]])
        v2 = np.array([p3[0] - p2[0], p3[1] - p2[1]])
        dot = np.dot(v1, v2)
        m1, m2 = np.linalg.norm(v1), np.linalg.norm(v2)
        if m1 * m2 == 0:
            return 180
        return math.degrees(math.acos(np.clip(dot / (m1 * m2), -1.0, 1.0)))

    def get_midpoint(self, p1, p2):
        return np.array([(p1[0] + p2[0]) / 2, (p1[1] + p2[1]) / 2])

    def update_pose_state(self, results):
        now = time.time()
        if not results or len(results) == 0:
            return self._handle_no_detection(now)
        result = results[0]
        if result.keypoints is None or len(result.keypoints) == 0:
            return self._handle_no_detection(now)

        kp = result.keypoints.xy.cpu().numpy()
        conf = result.keypoints.conf.cpu().numpy() if result.keypoints.conf is not None else None
        if kp.shape[0] == 0:
            return self._handle_no_detection(now)
        kp, conf = kp[0], conf[0] if conf is not None else np.ones(17)

        required = [self.L_SHOULDER, self.R_SHOULDER, self.L_HIP, self.R_HIP, self.L_KNEE, self.R_KNEE]
        if not all(conf[i] > 0.4 for i in required):
            return self._handle_no_detection(now)

        mid_shoulder = self.get_midpoint(kp[self.L_SHOULDER], kp[self.R_SHOULDER])
        mid_hip = self.get_midpoint(kp[self.L_HIP], kp[self.R_HIP])
        mid_knee = self.get_midpoint(kp[self.L_KNEE], kp[self.R_KNEE])

        is_kneeling = mid_knee[1] > mid_hip[1] + self.MIN_KNEE_HIP_DIFF
        spine_angle = self.calculate_angle(mid_shoulder, mid_hip, mid_knee)
        backbend_deviation = 180 - spine_angle

        if backbend_deviation >= self.EXCELLENT_BACKBEND_ANGLE:
            depth_score = 100
        elif backbend_deviation >= self.GOOD_BACKBEND_ANGLE:
            depth_score = 60 + (backbend_deviation - self.GOOD_BACKBEND_ANGLE) / (self.EXCELLENT_BACKBEND_ANGLE - self.GOOD_BACKBEND_ANGLE) * 40
        elif backbend_deviation >= self.MIN_BACKBEND_ANGLE:
            depth_score = 30 + (backbend_deviation - self.MIN_BACKBEND_ANGLE) / (self.GOOD_BACKBEND_ANGLE - self.MIN_BACKBEND_ANGLE) * 30
        else:
            depth_score = (backbend_deviation / self.MIN_BACKBEND_ANGLE) * 30

        self.depth_buffer.append(depth_score)
        self.current_backbend_depth = np.median(self.depth_buffer)

        valid = is_kneeling and backbend_deviation >= self.MIN_BACKBEND_ANGLE
        if valid:
            if not self.in_pose:
                self.in_pose = True
                self.pose_start_time = now
                self.current_state = "HOLDING"
            self.last_valid_time = now
            self.best_backbend_depth = max(self.best_backbend_depth, self.current_backbend_depth)
            self.current_hold_time = now - self.pose_start_time
            self.best_hold_time = max(self.best_hold_time, self.current_hold_time)
        else:
            self._check_grace_period(now)

    def _handle_no_detection(self, now):
        self._check_grace_period(now)

    def _check_grace_period(self, now):
        if self.in_pose and self.last_valid_time is not None:
            if now - self.last_valid_time > self.GRACE_SECONDS:
                self.total_hold_time += self.current_hold_time
                self.in_pose = False
                self.pose_start_time = None
                self.current_hold_time = 0.0
                self.current_state = "READY"

    def visualize(self, results):
        if not results or len(results) == 0:
            return np.zeros((480, 640, 3), dtype=np.uint8)
        frame = results[0].orig_img.copy()
        self.update_pose_state(results)
        return frame

    def get_current_stats(self):
        return {
            "current_time": round(self.current_hold_time, 1),
            "best_time": round(self.best_hold_time, 1),
            "total_time": round(self.total_hold_time + self.current_hold_time, 1),
            "state": self.current_state,
        }

    def reset(self):
        self.in_pose = False
        self.pose_start_time = None
        self.current_hold_time = 0.0
        self.best_hold_time = 0.0
        self.total_hold_time = 0.0
        self.current_backbend_depth = 0.0
        self.best_backbend_depth = 0.0
        self.current_state = "READY"
        self.depth_buffer.clear()


# ============== NATARAJASANA (DANCER POSE) MONITOR ==============
class NatarajasanaMonitor:
    def __init__(self, model_path=MODEL_PATH):
        self.model = YOLO(model_path)
        self.model.to(DEVICE)

        self.L_SHOULDER, self.R_SHOULDER = 5, 6
        self.L_HIP, self.R_HIP = 11, 12
        self.L_KNEE, self.R_KNEE = 13, 14
        self.L_ANKLE, self.R_ANKLE = 15, 16
        self.L_WRIST, self.R_WRIST = 9, 10

        self.in_pose = False
        self.pose_start_time = None
        self.current_hold_time = 0.0
        self.best_hold_time = 0.0
        self.total_hold_time = 0.0
        self.current_leg_height = 0.0
        self.best_leg_height = 0.0
        self.balance_score = 100.0
        self.overall_form_score = 0.0
        self.standing_foot_history = deque(maxlen=30)

        self.MIN_LEG_LIFT_ANGLE = 25
        self.GOOD_LEG_LIFT_ANGLE = 45
        self.EXCELLENT_LEG_LIFT_ANGLE = 70
        self.ONE_LEG_THRESHOLD = 80

        self.GRACE_SECONDS = 1.5
        self.last_valid_time = None
        self.height_buffer = deque(maxlen=5)
        self.form_buffer = deque(maxlen=5)
        self.current_state = "READY"
        self.detected_side = None

    def process_frame(self, frame):
        return self.model(frame, verbose=False, device=DEVICE)

    def calculate_angle(self, p1, p2, p3):
        v1 = np.array([p1[0] - p2[0], p1[1] - p2[1]])
        v2 = np.array([p3[0] - p2[0], p3[1] - p2[1]])
        dot = np.dot(v1, v2)
        m1, m2 = np.linalg.norm(v1), np.linalg.norm(v2)
        if m1 * m2 == 0:
            return 180
        return math.degrees(math.acos(np.clip(dot / (m1 * m2), -1.0, 1.0)))

    def _angle_from_vertical(self, p1, p2):
        dx = p2[0] - p1[0]
        dy = p2[1] - p1[1]
        return math.degrees(math.atan2(abs(dx), abs(dy)))

    def _detect_standing_side(self, kp, conf):
        if conf[self.L_ANKLE] < 0.4 or conf[self.R_ANKLE] < 0.4:
            return None
        l_y, r_y = kp[self.L_ANKLE][1], kp[self.R_ANKLE][1]
        if l_y > r_y + self.ONE_LEG_THRESHOLD:
            return "left"
        elif r_y > l_y + self.ONE_LEG_THRESHOLD:
            return "right"
        return None

    def _calc_balance(self, ankle):
        self.standing_foot_history.append(ankle.copy())
        if len(self.standing_foot_history) < 5:
            return 100.0
        positions = np.array(list(self.standing_foot_history))
        variance = np.var(positions, axis=0)
        return max(0, 100 - np.sqrt(variance[0] + variance[1]) * 2)

    def update_pose_state(self, results):
        now = time.time()
        if not results or len(results) == 0:
            return self._handle_no_detection(now)
        result = results[0]
        if result.keypoints is None or len(result.keypoints) == 0:
            return self._handle_no_detection(now)

        kp = result.keypoints.xy.cpu().numpy()
        conf = result.keypoints.conf.cpu().numpy() if result.keypoints.conf is not None else None
        if kp.shape[0] == 0:
            return self._handle_no_detection(now)
        kp, conf = kp[0], conf[0] if conf is not None else np.ones(17)

        side = self._detect_standing_side(kp, conf)
        if side is None:
            return self._handle_no_detection(now)
        self.detected_side = side

        if side == "left":
            s_hip, s_knee, s_ankle = self.L_HIP, self.L_KNEE, self.L_ANKLE
            l_hip, l_knee, l_ankle = self.R_HIP, self.R_KNEE, self.R_ANKLE
        else:
            s_hip, s_knee, s_ankle = self.R_HIP, self.R_KNEE, self.R_ANKLE
            l_hip, l_knee, l_ankle = self.L_HIP, self.L_KNEE, self.L_ANKLE

        required = [s_hip, s_knee, s_ankle, l_hip, l_knee, l_ankle]
        if not all(conf[i] > 0.4 for i in required):
            return self._handle_no_detection(now)

        leg_lift = self._angle_from_vertical(kp[l_hip], kp[l_ankle])
        standing_straight = abs(180 - self.calculate_angle(kp[s_hip], kp[s_knee], kp[s_ankle]))
        balance = self._calc_balance(kp[s_ankle])
        leg_height_pct = min(100, (leg_lift / self.EXCELLENT_LEG_LIFT_ANGLE) * 100)
        form_score = min(40, (leg_lift / self.EXCELLENT_LEG_LIFT_ANGLE) * 40) + max(0, 20 - standing_straight) + (balance / 100) * 25

        self.height_buffer.append(leg_height_pct)
        self.form_buffer.append(form_score)
        self.current_leg_height = np.median(self.height_buffer)
        self.balance_score = balance
        self.overall_form_score = np.median(self.form_buffer)

        valid = leg_lift >= self.MIN_LEG_LIFT_ANGLE and standing_straight <= 30
        if valid:
            if not self.in_pose:
                self.in_pose = True
                self.pose_start_time = now
                self.current_state = "HOLDING"
            self.last_valid_time = now
            self.best_leg_height = max(self.best_leg_height, self.current_leg_height)
            self.current_hold_time = now - self.pose_start_time
            self.best_hold_time = max(self.best_hold_time, self.current_hold_time)
        else:
            self._check_grace_period(now)

    def _handle_no_detection(self, now):
        self._check_grace_period(now)

    def _check_grace_period(self, now):
        if self.in_pose and self.last_valid_time is not None:
            if now - self.last_valid_time > self.GRACE_SECONDS:
                self.total_hold_time += self.current_hold_time
                self.in_pose = False
                self.pose_start_time = None
                self.current_hold_time = 0.0
                self.current_state = "READY"
                self.standing_foot_history.clear()

    def visualize(self, results):
        if not results or len(results) == 0:
            return np.zeros((480, 640, 3), dtype=np.uint8)
        frame = results[0].orig_img.copy()
        self.update_pose_state(results)
        return frame

    def get_current_stats(self):
        return {
            "current_time": round(self.current_hold_time, 1),
            "best_time": round(self.best_hold_time, 1),
            "total_time": round(self.total_hold_time + self.current_hold_time, 1),
            "state": self.current_state,
        }

    def reset(self):
        self.in_pose = False
        self.pose_start_time = None
        self.current_hold_time = 0.0
        self.best_hold_time = 0.0
        self.total_hold_time = 0.0
        self.current_leg_height = 0.0
        self.best_leg_height = 0.0
        self.balance_score = 100.0
        self.overall_form_score = 0.0
        self.current_state = "READY"
        self.detected_side = None
        self.standing_foot_history.clear()
        self.height_buffer.clear()
        self.form_buffer.clear()


# ============== PASCHIMOTTANASANA (SEATED FORWARD BEND) MONITOR ==============
class PaschimottanasanaMonitor:
    def __init__(self, model_path=MODEL_PATH):
        self.model = YOLO(model_path)
        self.model.to(DEVICE)

        self.L_SHOULDER, self.R_SHOULDER = 5, 6
        self.L_HIP, self.R_HIP = 11, 12
        self.L_KNEE, self.R_KNEE = 13, 14
        self.L_ANKLE, self.R_ANKLE = 15, 16

        self.in_pose = False
        self.pose_start_time = None
        self.current_hold_time = 0.0
        self.best_hold_time = 0.0
        self.total_hold_time = 0.0
        self.current_depth = 0.0
        self.best_depth = 0.0

        self.MIN_FORWARD_ANGLE = 20
        self.IDEAL_FORWARD_ANGLE = 45

        self.GRACE_SECONDS = 1.0
        self.last_valid_time = None
        self.depth_buffer = deque(maxlen=5)
        self.current_state = "READY"

    def process_frame(self, frame):
        return self.model(frame, verbose=False, device=DEVICE)

    def calculate_angle(self, p1, p2, p3):
        v1 = np.array([p1[0] - p2[0], p1[1] - p2[1]])
        v2 = np.array([p3[0] - p2[0], p3[1] - p2[1]])
        dot = np.dot(v1, v2)
        m1, m2 = np.linalg.norm(v1), np.linalg.norm(v2)
        if m1 * m2 == 0:
            return 180
        return math.degrees(math.acos(np.clip(dot / (m1 * m2), -1.0, 1.0)))

    def get_midpoint(self, p1, p2):
        return np.array([(p1[0] + p2[0]) / 2, (p1[1] + p2[1]) / 2])

    def update_pose_state(self, results):
        now = time.time()
        if not results or len(results) == 0:
            return self._handle_no_detection(now)
        result = results[0]
        if result.keypoints is None or len(result.keypoints) == 0:
            return self._handle_no_detection(now)

        kp = result.keypoints.xy.cpu().numpy()
        conf = result.keypoints.conf.cpu().numpy() if result.keypoints.conf is not None else None
        if kp.shape[0] == 0:
            return self._handle_no_detection(now)
        kp, conf = kp[0], conf[0] if conf is not None else np.ones(17)

        required = [self.L_SHOULDER, self.R_SHOULDER, self.L_HIP, self.R_HIP,
                    self.L_KNEE, self.R_KNEE, self.L_ANKLE, self.R_ANKLE]
        if not all(conf[i] > 0.4 for i in required):
            return self._handle_no_detection(now)

        mid_shoulder = self.get_midpoint(kp[self.L_SHOULDER], kp[self.R_SHOULDER])
        mid_hip = self.get_midpoint(kp[self.L_HIP], kp[self.R_HIP])
        mid_knee = self.get_midpoint(kp[self.L_KNEE], kp[self.R_KNEE])

        hip_angle = self.calculate_angle(mid_shoulder, mid_hip, mid_knee)
        forward_bend = 180 - hip_angle
        depth_pct = min(100, (forward_bend / 90) * 100)
        is_seated = mid_hip[1] > mid_shoulder[1] - 50

        self.depth_buffer.append(depth_pct)
        self.current_depth = np.median(self.depth_buffer)

        valid = forward_bend >= self.MIN_FORWARD_ANGLE and is_seated
        if valid:
            if not self.in_pose:
                self.in_pose = True
                self.pose_start_time = now
                self.current_state = "HOLDING"
            self.last_valid_time = now
            self.best_depth = max(self.best_depth, self.current_depth)
            self.current_hold_time = now - self.pose_start_time
            self.best_hold_time = max(self.best_hold_time, self.current_hold_time)
        else:
            self._check_grace_period(now)

    def _handle_no_detection(self, now):
        self._check_grace_period(now)

    def _check_grace_period(self, now):
        if self.in_pose and self.last_valid_time is not None:
            if now - self.last_valid_time > self.GRACE_SECONDS:
                self.total_hold_time += self.current_hold_time
                self.in_pose = False
                self.pose_start_time = None
                self.current_hold_time = 0.0
                self.current_state = "READY"

    def visualize(self, results):
        if not results or len(results) == 0:
            return np.zeros((480, 640, 3), dtype=np.uint8)
        frame = results[0].orig_img.copy()
        self.update_pose_state(results)
        return frame

    def get_current_stats(self):
        return {
            "current_time": round(self.current_hold_time, 1),
            "best_time": round(self.best_hold_time, 1),
            "total_time": round(self.total_hold_time + self.current_hold_time, 1),
            "state": self.current_state,
        }

    def reset(self):
        self.in_pose = False
        self.pose_start_time = None
        self.current_hold_time = 0.0
        self.best_hold_time = 0.0
        self.total_hold_time = 0.0
        self.current_depth = 0.0
        self.best_depth = 0.0
        self.current_state = "READY"
        self.depth_buffer.clear()


# ============== FASTAPI APP ==============
active_sessions: Dict[str, Dict] = {}

EXERCISE_MONITORS = {
    "pushups": PushupMonitor,
    "squats": SquatMonitor,
    "crunches": CrunchMonitor,
    "leg_raises": LegRaiseMonitor,
    "plank": PlankMonitor,
    "baddha_konasana": BaddhaKonasanaMonitor,
    "ustrasana": UstrasanaMonitor,
    "natarajasana": NatarajasanaMonitor,
    "paschimottanasana": PaschimottanasanaMonitor,
}

EXERCISE_TYPES = {
    "pushups": "reps",
    "squats": "reps",
    "crunches": "reps",
    "leg_raises": "reps",
    "plank": "timed",
    "baddha_konasana": "timed",
    "ustrasana": "timed",
    "natarajasana": "timed",
    "paschimottanasana": "timed",
}

@asynccontextmanager
async def lifespan(app: FastAPI):
    print(f"🚀 AuraFit AI started on device: {DEVICE}")
    yield
    print("👋 AuraFit AI shutting down...")
    active_sessions.clear()

app = FastAPI(title="AuraFit AI", version="1.0.0", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Serve static files (React build)
static_dir = Path("static")
if static_dir.exists():
    app.mount("/assets", StaticFiles(directory="static/assets"), name="assets")

@app.get("/")
async def root():
    """Serve React app"""
    index_path = Path("static/index.html")
    if index_path.exists():
        return FileResponse(index_path)
    return {"service": "AuraFit AI", "status": "active", "device": DEVICE}


@app.get("/api/health")
async def health():
    return {"status": "healthy", "device": DEVICE}


@app.get("/api/exercises")
async def get_exercises():
    return {
        "exercises": [
            {"id": "pushups", "name": "Pushups", "description": "Upper body strength", "icon": "💪", "type": "reps", "available": True},
            {"id": "squats", "name": "Squats", "description": "Lower body strength", "icon": "🦵", "type": "reps", "available": True},
            {"id": "crunches", "name": "Crunches", "description": "Core strength", "icon": "🔥", "type": "reps", "available": True},
            {"id": "leg_raises", "name": "Leg Raises", "description": "Lower abs", "icon": "🦿", "type": "reps", "available": True},
            {"id": "plank", "name": "Plank", "description": "Core endurance", "icon": "🧘", "type": "timed", "available": True},
            {"id": "baddha_konasana", "name": "Baddha Konasana", "description": "Hip flexibility (Butterfly Pose)", "icon": "🦋", "type": "timed", "available": True},
            {"id": "ustrasana", "name": "Ustrasana", "description": "Spine flexibility (Camel Pose)", "icon": "🐪", "type": "timed", "available": True},
            {"id": "natarajasana", "name": "Natarajasana", "description": "Balance & flexibility (Dancer Pose)", "icon": "💃", "type": "timed", "available": True},
            {"id": "paschimottanasana", "name": "Paschimottanasana", "description": "Hamstring flexibility (Seated Forward Bend)", "icon": "🧘‍♀️", "type": "timed", "available": True},
        ]
    }


@app.post("/api/session/start")
async def start_session(exercise: str, session_id: str):
    if exercise not in EXERCISE_MONITORS:
        raise HTTPException(status_code=400, detail=f"Exercise '{exercise}' not supported")
    
    if session_id in active_sessions:
        current_session = active_sessions[session_id]
        if current_session["exercise"] == exercise:
            return {
                "session_id": session_id,
                "exercise": current_session["exercise"],
                "status": "already_active",
                "message": f"Session already running for {exercise}"
            }
        else:
            del active_sessions[session_id]
    
    monitor = EXERCISE_MONITORS[exercise]()
    exercise_type = EXERCISE_TYPES.get(exercise, "reps")
    
    active_sessions[session_id] = {
        "exercise": exercise,
        "exercise_type": exercise_type,
        "monitor": monitor,
        "sets": 0,
        "active": True
    }
    
    return {
        "session_id": session_id,
        "exercise": exercise,
        "exercise_type": exercise_type,
        "status": "started",
        "message": f"Session started for {exercise}"
    }


@app.post("/api/session/stop")
async def stop_session(session_id: str):
    if session_id not in active_sessions:
        raise HTTPException(status_code=404, detail="Session not found")
    
    session = active_sessions[session_id]
    monitor = session["monitor"]
    exercise = session["exercise"]
    exercise_type = session["exercise_type"]
    
    if exercise_type == "timed":
        stats = monitor.get_current_stats()
        final_stats = {
            "session_id": session_id,
            "exercise": exercise,
            "exercise_type": "timed",
            "total_time": stats.get("total_time", 0),
            "best_time": stats.get("best_time", 0),
            "sets": session["sets"],
            "status": "completed"
        }
    else:
        total_reps = _get_rep_count(monitor)
        final_stats = {
            "session_id": session_id,
            "exercise": exercise,
            "exercise_type": "reps",
            "total_reps": total_reps,
            "sets": session["sets"],
            "status": "completed"
        }
    
    del active_sessions[session_id]
    return final_stats


def _get_rep_count(monitor):
    for attr in ['pushup_count', 'squat_count', 'crunch_count', 'rep_count']:
        if hasattr(monitor, attr):
            return getattr(monitor, attr)
    return 0


def _get_timed_stats(monitor):
    stats = monitor.get_current_stats()
    return {
        "current_time": stats.get("current_time", stats.get("current_elapsed", 0)),
        "best_time": stats.get("best_time", stats.get("best_duration", 0)),
        "total_time": stats.get("total_time", stats.get("total_plank_time", 0)),
        "state": stats.get("state", "UNKNOWN"),
    }


@app.get("/api/session/stats")
async def get_session_stats(session_id: str):
    if session_id not in active_sessions:
        raise HTTPException(status_code=404, detail="Session not found")
    
    session = active_sessions[session_id]
    monitor = session["monitor"]
    exercise = session["exercise"]
    exercise_type = session["exercise_type"]
    
    if exercise_type == "timed":
        timed = _get_timed_stats(monitor)
        return {
            "session_id": session_id,
            "exercise": exercise,
            "exercise_type": "timed",
            "current_time": timed["current_time"],
            "best_time": timed["best_time"],
            "total_time": timed["total_time"],
            "state": timed["state"],
            "sets": session["sets"],
            "active": session["active"]
        }
    else:
        reps = _get_rep_count(monitor)
        return {
            "session_id": session_id,
            "exercise": exercise,
            "exercise_type": "reps",
            "reps": reps,
            "sets": session["sets"],
            "state": monitor.current_state,
            "active": session["active"]
        }


def _build_ws_metadata(monitor, session, exercise_type):
    exercise = session["exercise"]
    if exercise_type == "timed":
        timed = _get_timed_stats(monitor)
        feedback = _get_timed_feedback(timed["state"], exercise)
        return {
            "type": "metadata",
            "exercise_type": "timed",
            "current_time": timed["current_time"],
            "best_time": timed["best_time"],
            "total_time": timed["total_time"],
            "sets": session["sets"],
            "state": timed["state"],
            "feedback": feedback
        }
    else:
        count = _get_rep_count(monitor)
        feedback = _get_rep_feedback(monitor.current_state, exercise)
        return {
            "type": "metadata",
            "exercise_type": "reps",
            "reps": count,
            "sets": session["sets"],
            "state": monitor.current_state,
            "feedback": feedback
        }


def _build_ws_response(monitor, session, exercise_type, frame_base64):
    exercise = session["exercise"]
    if exercise_type == "timed":
        timed = _get_timed_stats(monitor)
        feedback = _get_timed_feedback(timed["state"], exercise)
        return {
            "type": "processed_frame",
            "frame": f"image/jpeg;base64,{frame_base64}",
            "exercise_type": "timed",
            "current_time": timed["current_time"],
            "best_time": timed["best_time"],
            "total_time": timed["total_time"],
            "sets": session["sets"],
            "state": timed["state"],
            "feedback": feedback
        }
    else:
        count = _get_rep_count(monitor)
        feedback = _get_rep_feedback(monitor.current_state, exercise)
        return {
            "type": "processed_frame",
            "frame": f"image/jpeg;base64,{frame_base64}",
            "exercise_type": "reps",
            "reps": count,
            "sets": session["sets"],
            "state": monitor.current_state,
            "feedback": feedback
        }


def _get_rep_feedback(state: str, exercise: str) -> list:
    exercise_tips = {
        "pushups": {"UP": "Great form! Keep your core engaged.", "DOWN": "Control the descent, chest to floor."},
        "squats": {"UP": "Drive through your heels!", "DOWN": "Keep your knees over toes."},
        "crunches": {"UP": "Squeeze your abs at the top!", "DOWN": "Control the lowering motion."},
        "leg_raises": {"UP": "Keep your legs straight!", "DOWN": "Don't let your feet touch the ground."},
    }
    tips = exercise_tips.get(exercise, {"UP": "Good!", "DOWN": "Keep going!"})
    if state == "UP":
        return [tips["UP"]]
    elif state == "DOWN":
        return [tips["DOWN"]]
    return ["Position yourself in the frame."]


def _get_timed_feedback(state: str, exercise: str) -> list:
    timed_tips = {
        "plank": {"active": "Great form! Keep holding!", "rest": "Get into plank position - hands under shoulders."},
        "baddha_konasana": {"active": "Relax into the stretch, breathe deeply.", "rest": "Sit with soles of feet together, knees out."},
        "ustrasana": {"active": "Open your chest, breathe!", "rest": "Kneel and lean back into camel pose."},
        "natarajasana": {"active": "Beautiful balance! Hold steady.", "rest": "Stand on one leg, grab back foot."},
        "paschimottanasana": {"active": "Fold deeper with each exhale.", "rest": "Sit with legs extended, fold forward."},
    }
    tips = timed_tips.get(exercise, {"active": "Hold the pose!", "rest": "Get into position."})
    active_states = {"PLANK", "HOLDING", "IN_POSE"}
    if state in active_states:
        return [tips["active"]]
    return [tips["rest"]]


@app.websocket("/ws/exercise/{session_id}")
async def websocket_endpoint(websocket: WebSocket, session_id: str):
    await websocket.accept()
    
    if session_id not in active_sessions:
        await websocket.send_json({"error": "Session not found. Please start a session first."})
        await websocket.close()
        return
    
    session = active_sessions[session_id]
    monitor = session["monitor"]
    exercise = session["exercise"]
    exercise_type = session.get("exercise_type", EXERCISE_TYPES.get(exercise, "reps"))
    
    binary_mode = websocket.query_params.get("mode", "json") == "binary"
    
    try:
        while True:
            data = await websocket.receive_text()
            message = json.loads(data)
            
            if message.get("type") == "frame":
                try:
                    frame_data = message["frame"]
                    if "," in frame_data:
                        frame_data = frame_data.split(",")[1]
                    
                    img_bytes = base64.b64decode(frame_data)
                    nparr = np.frombuffer(img_bytes, np.uint8)
                    frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
                    
                    if frame is None:
                        continue
                    
                    frame = _resize_frame(frame)
                    results = monitor.process_frame(frame)
                    annotated = monitor.visualize(results)
                    
                    if annotated is None:
                        annotated = frame
                    
                    if binary_mode:
                        _, buffer = cv2.imencode('.jpg', annotated,
                                                 [cv2.IMWRITE_JPEG_QUALITY, JPEG_QUALITY])
                        await websocket.send_bytes(buffer.tobytes())
                        metadata = _build_ws_metadata(monitor, session, exercise_type)
                        await websocket.send_json(metadata)
                    else:
                        _, buffer = cv2.imencode('.jpg', annotated,
                                                 [cv2.IMWRITE_JPEG_QUALITY, JPEG_QUALITY])
                        frame_b64 = base64.b64encode(buffer).decode('utf-8')
                        response = _build_ws_response(monitor, session, exercise_type, frame_b64)
                        await websocket.send_json(response)
                
                except Exception as e:
                    print(f"Frame error: {e}")
            
            elif message.get("type") == "complete_set":
                session["sets"] += 1
                if hasattr(monitor, 'reset_counter'):
                    monitor.reset_counter()
                elif hasattr(monitor, 'reset'):
                    monitor.reset()
                await websocket.send_json({
                    "type": "set_completed",
                    "sets": session["sets"],
                    "message": f"Set {session['sets']} completed!"
                })
            
            elif message.get("type") == "ping":
                await websocket.send_json({"type": "pong"})
    
    except WebSocketDisconnect:
        print(f"Disconnected: {session_id}")
        session["active"] = False
    except Exception as e:
        print(f"WS Error: {e}")


# Catch-all for React Router
@app.get("/{full_path:path}")
async def serve_spa(full_path: str):
    """Serve React app for all other routes"""
    # Check if it's a static file request
    static_file = Path(f"static/{full_path}")
    if static_file.exists() and static_file.is_file():
        return FileResponse(static_file)
    # Otherwise serve index.html for SPA routing
    index_path = Path("static/index.html")
    if index_path.exists():
        return FileResponse(index_path)
    return {"error": "Not found"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")