# app.py
import os
import sys # Import sys to print to stderr for logging
import traceback # Import traceback for detailed error printing

# <<< Force disable Numba cache BEFORE any imports that might use it >>>
# Printing to stderr ensures it shows up in container logs easily
print("----------------------------------------------------------", file=sys.stderr)
print("Attempting to set NUMBA_DISABLE_CACHE=1", file=sys.stderr)
os.environ['NUMBA_DISABLE_CACHE'] = '1'
# <<< Verify it was set (check logs for this output) >>>
numba_cache_status = os.getenv('NUMBA_DISABLE_CACHE')
print(f"NUMBA_DISABLE_CACHE is now set to: {numba_cache_status}", file=sys.stderr)
if numba_cache_status != '1':
    print("WARNING: NUMBA_DISABLE_CACHE was not set correctly!", file=sys.stderr)
print("----------------------------------------------------------", file=sys.stderr)

# Now proceed with other imports
print("Importing Flask, base64...", file=sys.stderr)
from flask import Flask, render_template, request, jsonify
import base64

# Try importing rembg with detailed logging
try:
    print("Attempting to import rembg...", file=sys.stderr)
    # Flush output buffer to ensure messages appear in order
    sys.stderr.flush()
    from rembg import remove  # Import the remove function from rembg
    print("Imported rembg successfully.", file=sys.stderr)
    sys.stderr.flush()
except ImportError as e:
     # Catch ImportError specifically if rembg isn't found
     print(f"FATAL: Could not import rembg: {e}", file=sys.stderr)
     print("Please check installation and dependencies in requirements.txt.", file=sys.stderr)
     sys.stderr.flush()
     raise # Re-raise to stop execution clearly
except Exception as e:
    # Catch other potential exceptions during import (like the Numba error)
    print(f"ERROR during import process (likely rembg or its dependency): {e}", file=sys.stderr)
    # Print the full traceback for detailed debugging
    traceback.print_exc(file=sys.stderr)
    sys.stderr.flush()
    raise # Re-raise the exception to see the full traceback in logs

print("Flask app initialization starting...", file=sys.stderr)
app = Flask(__name__)
print("Flask app initialized.", file=sys.stderr)
sys.stderr.flush()

@app.route('/')
def index():
    # Render the main HTML page
    print("Serving index page.", file=sys.stderr)
    sys.stderr.flush()
    return render_template('index.html')

@app.route('/upload', methods=['POST'])
def upload_file():
    print("Received request to /upload", file=sys.stderr)
    sys.stderr.flush()
    # Check if an image file is uploaded in the request
    if 'image' not in request.files:
        print("Error: No image part in the request.", file=sys.stderr)
        sys.stderr.flush()
        return jsonify({'error': 'No image part in the request'}), 400

    uploaded_file = request.files['image']

    # Check if the filename is empty (no file selected)
    if uploaded_file.filename == '':
        print("Error: No image selected (filename empty).", file=sys.stderr)
        sys.stderr.flush()
        return jsonify({'error': 'No image selected'}), 400

    # Check if the file object itself exists
    if not uploaded_file:
        print("Error: Uploaded file object is invalid.", file=sys.stderr)
        sys.stderr.flush()
        return jsonify({'error': 'Image file is invalid'}), 400

    print(f"Processing uploaded file: {uploaded_file.filename}", file=sys.stderr)
    sys.stderr.flush()
    try:
        # Read the uploaded image file bytes
        input_image_bytes = uploaded_file.read()
        print(f"Read {len(input_image_bytes)} bytes from image.", file=sys.stderr)
        sys.stderr.flush()

        # Use rembg to remove the background
        print("Calling rembg.remove...", file=sys.stderr)
        sys.stderr.flush()
        output_image_bytes = remove(input_image_bytes)
        print(f"rembg.remove finished, output size: {len(output_image_bytes)} bytes.", file=sys.stderr)
        sys.stderr.flush()

        # Convert the original image and background-removed image to base64
        print("Encoding images to base64...", file=sys.stderr)
        sys.stderr.flush()
        original_image_base64 = base64.b64encode(input_image_bytes).decode('utf-8')
        modified_image_base64 = base64.b64encode(output_image_bytes).decode('utf-8')
        print("Encoding complete.", file=sys.stderr)
        sys.stderr.flush()

        # Return the images as JSON to the frontend
        # Ensure the data URI format is correct (it needs image type, e.g., image/png)
        # rembg typically outputs PNG
        print("Returning JSON response.", file=sys.stderr)
        sys.stderr.flush()
        return jsonify({
            'original_image': f'data:image/png;base64,{original_image_base64}',
            'background_removed_image': f'data:image/png;base64,{modified_image_base64}'
        })

    except Exception as e:
        # Log the exception for debugging
        print(f"ERROR processing image with rembg: {str(e)}", file=sys.stderr)
        print("------------------- Traceback -------------------", file=sys.stderr)
        traceback.print_exc(file=sys.stderr)
        print("-------------------------------------------------", file=sys.stderr)
        sys.stderr.flush()
        # Return a generic error message to the user
        return jsonify({'error': f'Failed to process image: {str(e)}'}), 500

# Keep the __main__ block for running locally, but it's not used by Gunicorn
if __name__ == '__main__':
    print("Running Flask app directly (for local testing only)...", file=sys.stderr)
    # When run directly, use debug=True for development ease
    # When run with Gunicorn (like in Hugging Face), Gunicorn handles host/port/workers, debug should be False.
    app.run(host='0.0.0.0', port=5000, debug=True)