main.py

import os
import tempfile
import subprocess
from pathlib import Path
from flask import Flask, request, jsonify
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)

# Configure maximum file size (16MB)
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024

@app.route('/itl', methods=['POST'])
def image_to_latex():
    """Convert uploaded image to LaTeX code using TextTeller."""
    try:
        # Check if image file is present
        if 'image' not in request.files:
            return jsonify({'error': 'No image file provided'}), 400
        
        file = request.files['image']
        if file.filename == '':
            return jsonify({'error': 'No file selected'}), 400
        
        # Validate file type
        allowed_extensions = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'tiff'}
        file_ext = file.filename.rsplit('.', 1)[-1].lower() if '.' in file.filename else ''
        if file_ext not in allowed_extensions:
            return jsonify({'error': f'Invalid file type. Allowed: {", ".join(allowed_extensions)}'}), 400
        
        # Create temporary file to save uploaded image
        with tempfile.NamedTemporaryFile(delete=False, suffix=f'.{file_ext}') as tmp_file:
            file.save(tmp_file.name)
            temp_image_path = tmp_file.name
        
        try:
            # Run texteller inference command
            logger.info(f"Processing image: {temp_image_path}")
            result = subprocess.run(
                ['texteller', 'inference', temp_image_path],
                capture_output=True,
                text=True,
                timeout=120  # Increased to 2 minutes
            )
            
            if result.returncode == 0:
                # Extract LaTeX from output
                latex_output = result.stdout.strip()
                logger.info(f"Successfully processed image. LaTeX length: {len(latex_output)}")
                return jsonify({
                    'success': True,
                    'latex': latex_output
                })
            else:
                logger.error(f"TextTeller inference failed with return code {result.returncode}")
                logger.error(f"STDOUT: {result.stdout}")
                logger.error(f"STDERR: {result.stderr}")
                return jsonify({
                    'error': 'Failed to process image',
                    'details': result.stderr,
                    'return_code': result.returncode
                }), 500
        
        except subprocess.TimeoutExpired:
            logger.error("TextTeller inference timed out")
            return jsonify({'error': 'Processing timed out'}), 408
        
        except Exception as e:
            logger.error(f"Error during processing: {str(e)}")
            return jsonify({'error': f'Processing error: {str(e)}'}), 500
        
        finally:
            # Clean up temporary file
            try:
                os.unlink(temp_image_path)
                logger.info(f"Cleaned up temporary file: {temp_image_path}")
            except OSError as e:
                logger.warning(f"Failed to delete temporary file {temp_image_path}: {e}")
    
    except Exception as e:
        logger.error(f"Unexpected error: {str(e)}")
        return jsonify({'error': f'Unexpected error: {str(e)}'}), 500

@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint."""
    try:
        # Test if texteller is available
        result = subprocess.run(['texteller', '--help'], capture_output=True, timeout=10)
        texteller_available = result.returncode == 0
        
        # Also check version if available
        version_info = ""
        if texteller_available:
            try:
                version_result = subprocess.run(['texteller', '--version'], capture_output=True, timeout=5, text=True)
                if version_result.returncode == 0:
                    version_info = version_result.stdout.strip()
            except:
                version_info = "Version check failed"
        
        return jsonify({
            'status': 'healthy' if texteller_available else 'unhealthy',
            'texteller_available': texteller_available,
            'texteller_version': version_info,
            'texteller_help_output': result.stdout.decode('utf-8') if texteller_available else None
        })
    except Exception as e:
        return jsonify({
            'status': 'unhealthy',
            'error': str(e)
        }), 503

@app.route('/', methods=['GET'])
def root():
    """Root endpoint with API documentation."""
    return jsonify({
        'service': 'Image to LaTeX API',
        'version': '1.0.0',
        'endpoints': {
            'POST /itl': 'Convert image to LaTeX. Send image file as multipart/form-data with key "image"',
            'GET /health': 'Health check endpoint',
            'GET /': 'This documentation'
        },
        'supported_formats': ['png', 'jpg', 'jpeg', 'gif', 'bmp', 'tiff'],
        'max_file_size': '16MB'
    })

if __name__ == '__main__':
    # Check if texteller is installed
    try:
        result = subprocess.run(['texteller', '--help'], capture_output=True)
        if result.returncode != 0:
            logger.warning("TextTeller might not be properly installed")
    except FileNotFoundError:
        logger.error("TextTeller is not installed. Please install it with: pip install texteller")
    
    app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 7860)), debug=False)