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)

def get_yolo_boxes(filename):
    name, _ = os.path.splitext(filename)
    label_path = os.path.join(LABELS_DIR, name + '.txt')
    boxes = []
    if os.path.exists(label_path):
        with open(label_path, 'r') as f:
            for line in f:
                parts = line.strip().split()
                if len(parts) >= 5:
                    # format: class x_center y_center width height
                    try:
                        boxes.append({
                            'class': parts[0],
                            'x': float(parts[1]),
                            'y': float(parts[2]),
                            'w': float(parts[3]),
                            'h': float(parts[4])
                        })
                    except ValueError:
                        continue
    return boxes

@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_index = None
    next_image = None

    for i, img in enumerate(images):
        if is_labeled(img):
            labeled_count += 1
        elif next_image is None:
            next_image = img
            next_index = i
            
    # If all labeled, next_image remains None
    
    return jsonify({
        'total': total,
        'labeled_count': labeled_count,
        'index': next_index if next_index is not None else -1,
        'filename': next_image,
        'completed': next_image is None,
        'boxes': get_yolo_boxes(next_image) if next_image else []
    })

@app.route('/api/image/<int:index>')
def get_image_by_index(index):
    images = get_image_files()
    total = len(images)
    
    if total == 0:
        return jsonify({'error': 'No images'}), 404
        
    # Clamp index
    if index < 0: index = 0
    if index >= total: index = total - 1
    
    filename = images[index]
    
    # Calculate stats
    labeled_count = sum(1 for img in images if is_labeled(img))
    
    return jsonify({
        'total': total,
        'labeled_count': labeled_count,
        'index': index,
        'filename': filename,
        'labeled': is_labeled(filename),
        'boxes': get_yolo_boxes(filename)
    })

@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)