#!/usr/bin/env python3
"""
PaddleOCR Server v5 (PP-OCRv5)
Flask HTTP server exposing PaddleOCR v3.3.0 functionality
"""

from paddlex import create_model
import base64
import numpy as np
from PIL import Image
from io import BytesIO
from flask import Flask, request, jsonify
import traceback

app = Flask(__name__)

# Initialize PP-OCRv5 model
print("Initializing PP-OCRv5 model...")
model = create_model("PP-OCRv5_server")
print("PP-OCRv5 model loaded successfully!")

@app.route('/health', methods=['GET'])
def health():
    """Health check endpoint."""
    return jsonify({
        'status': 'ok',
        'service': 'paddleocr-server-v5',
        'version': '3.3.0',
        'model': 'PP-OCRv5_server',
        'gpu_enabled': True
    })

@app.route('/ocr', methods=['POST'])
def ocr_endpoint():
    """
    OCR endpoint using PP-OCRv5.

    Accepts: {"image": "base64_encoded_image"}
    Returns: {"success": true, "count": N, "results": [...]}
    """
    try:
        # Parse request
        data = request.get_json()
        image_base64 = data['image']

        # Decode image
        image_bytes = base64.b64decode(image_base64)
        image = Image.open(BytesIO(image_bytes))
        image_np = np.array(image)

        # Run OCR with PP-OCRv5
        result = model.predict(image_np)

        # Format results
        formatted_results = []

        if result and 'dt_polys' in result[0] and 'rec_text' in result[0]:
            dt_polys = result[0]['dt_polys']
            rec_texts = result[0]['rec_text']
            rec_scores = result[0]['rec_score']

            for i in range(len(dt_polys)):
                box = dt_polys[i].tolist()  # Convert to list
                text = rec_texts[i]
                confidence = float(rec_scores[i])

                formatted_results.append({
                    'box': box,
                    'text': text,
                    'confidence': confidence
                })

        return jsonify({
            'success': True,
            'count': len(formatted_results),
            'results': formatted_results
        })

    except Exception as e:
        print(f"Error during OCR: {str(e)}")
        traceback.print_exc()
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500

if __name__ == '__main__':
    print("Starting PP-OCRv5 server on port 5555...")
    print("Model: PP-OCRv5_server")
    print("Version: 3.3.0")
    app.run(host='0.0.0.0', port=5555, debug=False)