import os
from flask import Flask, render_template, send_from_directory, jsonify, request

app = Flask(__name__)

IMAGES_DIR = 'images'
LABELS_DIR = 'labels'
os.makedirs(LABELS_DIR, exist_ok=True)

ALLOWED_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.bmp', '.webp'}

def get_image_files():
    files = []
    if os.path.exists(IMAGES_DIR):
        for f in os.listdir(IMAGES_DIR):
            _, ext = os.path.splitext(f)
            if ext.lower() in ALLOWED_EXTENSIONS:
                files.append(f)
    return sorted(files)

def is_labeled(filename):
    name, _ = os.path.splitext(filename)
    label_path = os.path.join(LABELS_DIR, name + '.txt')
    return os.path.exists(label_path)

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

@app.route('/images/<path:filename>')
def serve_image(filename):
    return send_from_directory(IMAGES_DIR, filename)

@app.route('/api/next')
def get_next_image():
    images = get_image_files()
    total = len(images)
    labeled_count = 0
    next_image = None

    for img in images:
        if is_labeled(img):
            labeled_count += 1
        elif next_image is None:
            next_image = img
            
    # If all labeled, next_image remains None
    
    return jsonify({
        'total': total,
        'labeled': labeled_count,
        'filename': next_image,
        'completed': next_image is None
    })

@app.route('/api/save', methods=['POST'])
def save_label():
    data = request.json
    filename = data.get('filename')
    boxes = data.get('boxes', []) # List of {x, y, w, h} in image pixels
    img_w = data.get('width')
    img_h = data.get('height')
    
    if not filename or not img_w or not img_h:
        return jsonify({'error': 'Missing data'}), 400

    name, _ = os.path.splitext(filename)
    label_path = os.path.join(LABELS_DIR, name + '.txt')
    
    lines = []
    for box in boxes:
        # YOLO format: class x_center y_center width height (normalized 0-1)
        # Class is 0 for signature
        
        # Ensure box is valid
        x = float(box['x'])
        y = float(box['y'])
        w = float(box['w'])
        h = float(box['h'])
        
        # Center coordinates
        center_x = (x + w / 2.0) / img_w
        center_y = (y + h / 2.0) / img_h
        norm_w = w / img_w
        norm_h = h / img_h
        
        # Clamp values just in case
        center_x = max(0.0, min(1.0, center_x))
        center_y = max(0.0, min(1.0, center_y))
        norm_w = max(0.0, min(1.0, norm_w))
        norm_h = max(0.0, min(1.0, norm_h))
        
        lines.append(f"0 {center_x:.6f} {center_y:.6f} {norm_w:.6f} {norm_h:.6f}")
        
    with open(label_path, 'w') as f:
        f.write('\n'.join(lines))
        
    return jsonify({'success': True})

if __name__ == '__main__':
    app.run(debug=True, port=8080)