from flask import Flask, request, jsonify
from ultralytics import YOLO
import cv2
import numpy as np
import os
import sys

# Initialize Flask App
app = Flask(__name__)

# Load YOLOv8 Model
# Check for custom user model in ai_runtime
CUSTOM_MODEL_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'ai_runtime', 'model.pt')
DEFAULT_MODEL_PATH = os.path.join(os.path.dirname(__file__), 'models', 'yolov8n.pt')

if os.path.exists(CUSTOM_MODEL_PATH):
    print(f"Loading Custom Model: {CUSTOM_MODEL_PATH}")
    model = YOLO(CUSTOM_MODEL_PATH)
elif os.path.exists(DEFAULT_MODEL_PATH):
    print(f"Loading Default Model: {DEFAULT_MODEL_PATH}")
    model = YOLO(DEFAULT_MODEL_PATH)
else:
    print("Downloading/Loading standard yolov8n.pt...")
    print("Downloading/Loading standard yolov8n.pt...")
    model = YOLO('yolov8n.pt')

print("Using Model Version: Rollback (train/weights/*.pt)")

@app.route('/detect', methods=['POST'])
def detect():
    try:
        if 'image' not in request.files:
            return jsonify({"status": "error", "message": "No image file provided"}), 400

        file = request.files['image']
        # Default confidence lowered to 30% to improve detection rate
        confidence = float(request.form.get('confidence', 30)) / 100.0

        # Read image
        file_bytes = np.frombuffer(file.read(), np.uint8)
        img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)

        if img is None:
            return jsonify({"status": "error", "message": "Failed to decode image"}), 400

        # Perform Inference (Fixed Resolution for Speed & Consistency)
        results = model.predict(img, conf=confidence, classes=[0], imgsz=640, verbose=False) # class 0 = person

        # Count persons
        count = 0
        for result in results:
            count += len(result.boxes)

        return jsonify({
            "status": "success",
            "student_count": count,
            "message": "Detection successful"
        })

    except Exception as e:
        import traceback
        traceback.print_exc()
        return jsonify({"status": "error", "message": f"Server Error: {str(e)}"}), 500

if __name__ == '__main__':
    print("Starting ClassVision AI Server on Port 5000 (DEBUG MODE)...")
    app.run(host='127.0.0.1', port=5000, debug=True, use_reloader=False)