from flask import Blueprint, request, jsonify, send_file, render_template
from app.utils.detector import MemoryDetector
import os
from PIL import Image
import io

main_bp = Blueprint('main', __name__)
detector = MemoryDetector()

@main_bp.route('/')
def index():
    return render_template('index.html')

@main_bp.route('/detect', methods=['POST'])
def detect_memory():
    if 'image' not in request.files:
        return jsonify({'error': 'No image provided'}), 400
    
    file = request.files['image']
    
    # Read the image
    img = Image.open(file.stream)
    
    # Process the image and get annotated image and detections
    annotated_img, detections = detector.detect(img)
    
    # Convert PIL image to bytes
    img_byte_arr = io.BytesIO()
    annotated_img.save(img_byte_arr, format='PNG')
    img_byte_arr.seek(0)
    
    return send_file(
        img_byte_arr,
        mimetype='image/png'
    )

@main_bp.route('/detect/test', methods=['GET'])
def detect_test():
    """Endpoint for testing with a hardcoded image"""
    # Using an existing image from the validation set
    test_image_path = os.path.join('training', 'val', 'images', 'memory_out19.png')
    
    if not os.path.exists(test_image_path):
        return jsonify({'error': f'Test image not found at {test_image_path}'}), 404
    
    img = Image.open(test_image_path)
    # Get both the annotated image and detections
    annotated_img, detections = detector.detect(img)
    
    img_byte_arr = io.BytesIO()
    annotated_img.save(img_byte_arr, format='PNG')
    img_byte_arr.seek(0)
    
    return send_file(
        img_byte_arr,
        mimetype='image/png'
    )