import os
import io
import time
import cv2
import numpy as np
from flask import Flask, render_template, Response, request, send_file, jsonify
from ultralytics import YOLO
from PIL import Image

app = Flask(__name__, template_folder="templates")

# ---------- CONFIG ----------
MODEL_PATH = "./bestbest.pt"      # ensure model is present
IMGSZ      = 540
WARNING_PX = 100
BOX_W = 100
BOX_H = 100
# ----------------------------

# 1. SETUP FACE DETECTION (Standard OpenCV Haarcascade)
# This loads the face detection model built into OpenCV
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

# Load Hand model once
model = None
if os.path.exists(MODEL_PATH):
    try:
        model = YOLO(MODEL_PATH)
        print("Model loaded.")
    except Exception as e:
        print("Model load failed:", e)
        model = None
else:
    print(f"WARNING: model not found at {MODEL_PATH}")

# --- helper functions (exact logic) ---
def get_box_coords(frame_w, frame_h, box_w=BOX_W, box_h=BOX_H):
    cx, cy = frame_w // 2, frame_h // 2
    x1 = cx - box_w // 2
    y1 = cy - box_h // 2
    x2 = cx + box_w // 2
    y2 = cy + box_h // 2
    return int(x1), int(y1), int(x2), int(y2)

def rects_intersect(rectA, rectB):
    ax1, ay1, ax2, ay2 = rectA
    bx1, by1, bx2, by2 = rectB
    inter_x1 = max(ax1, bx1)
    inter_y1 = max(ay1, by1)
    inter_x2 = min(ax2, bx2)
    inter_y2 = min(ay2, by2)
    iw = inter_x2 - inter_x1
    ih = inter_y2 - inter_y1
    return (iw > 0) and (ih > 0)

def rect_distance(rectA, rectB):
    ax1, ay1, ax2, ay2 = rectA
    bx1, by1, bx2, by2 = rectB
    dx = 0
    if ax2 < bx1:      dx = bx1 - ax2
    elif bx2 < ax1:    dx = ax1 - bx2
    dy = 0
    if ay2 < by1:      dy = by1 - ay2
    elif by2 < ay1:    dy = ay1 - by2
    return int(np.hypot(dx, dy))

def annotate_frame_bgr(frame_bgr):
    """Apply your detection + drawing logic on a BGR OpenCV frame, return annotated BGR"""
    h, w = frame_bgr.shape[:2]
    x1, y1, x2, y2 = get_box_coords(w, h, BOX_W, BOX_H)
    cv2.rectangle(frame_bgr, (x1, y1), (x2, y2), (220,220,220), 2)
    cv2.putText(frame_bgr, "DANGER ZONE", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (220,220,220), 2)

    # --- STEP 1: Detect Hand (Your original logic) ---
    hand_bbox = None
    if model is not None:
        try:
            results = model.predict(frame_bgr, imgsz=IMGSZ, conf=0.25, max_det=1, verbose=False)
            if len(results) > 0 and len(results[0].boxes) > 0:
                b = results[0].boxes[0]
                xyxy = b.xyxy.cpu().numpy().astype(int)[0]
                hand_bbox = (int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3]))
        except Exception as e:
            print("Inference error:", e)
            hand_bbox = None

    # --- STEP 2: Detect Faces & Filter ---
    # Convert to grayscale for face detection (required by haar cascade)
    gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray, 1.1, 4)

    # Check if the detected hand overlaps with ANY detected face
    if hand_bbox is not None:
        for (fx, fy, fw, fh) in faces:
            # Create face rectangle (x1, y1, x2, y2)
            face_rect = (fx, fy, fx + fw, fy + fh)
            
            # If hand overlaps face, HIDE HAND (set to None)
            if rects_intersect(hand_bbox, face_rect):
                print("Ignored hand detection (Overlap with Face)")
                hand_bbox = None
                break # Stop checking other faces, we already hid the hand

    # --- STEP 3: Draw Logic (Only draws if hand_bbox is still valid) ---
    state = "NO HAND"
    color = (100,100,100)
    dist = 0

    if hand_bbox is not None:
        danger_rect = (x1, y1, x2, y2)
        if rects_intersect(hand_bbox, danger_rect):
            state = "DANGER"; color = (0,0,255); dist = 0
        else:
            dist = rect_distance(hand_bbox, danger_rect)
            if dist <= WARNING_PX:
                state = "WARNING"; color = (0,165,255)
            else:
                state = "SAFE"; color = (0,255,0)

        cv2.rectangle(frame_bgr, (hand_bbox[0], hand_bbox[1]), (hand_bbox[2], hand_bbox[3]), color, 2)
        cv2.putText(frame_bgr, f"Dist: {dist}px", (10,110), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)

    cv2.putText(frame_bgr, state, (10,70), cv2.FONT_HERSHEY_DUPLEX, 1.5, color, 3)
    if state == "DANGER":
        cv2.putText(frame_bgr, "HAND TOO CLOSE!", (w//2 - 200, h//2), cv2.FONT_HERSHEY_DUPLEX, 1.5, (0,0,255), 4)

    return frame_bgr

# ---- Routes ----
@app.route("/")
def index():
    return render_template("index.html")

@app.route("/predict", methods=["POST"])
def predict():
    """
    Receives a single JPEG (multipart form under key 'frame').
    Returns processed JPEG image as response (image/jpeg).
    """
    if "frame" not in request.files:
        return jsonify({"error":"no frame"}), 400
    file = request.files["frame"]
    in_bytes = file.read()
    # read into OpenCV
    nparr = np.frombuffer(in_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)  # BGR
    if img is None:
        return jsonify({"error":"invalid image"}), 400

    # annotate
    out_bgr = annotate_frame_bgr(img)

    # encode to jpeg
    ret, buf = cv2.imencode(".jpg", out_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), 85])
    if not ret:
        return jsonify({"error":"encode failed"}), 500
    return Response(buf.tobytes(), mimetype="image/jpeg")

if __name__ == "__main__":
    # Run with flask for debugging
    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 7860)), debug=False)