import torch
from PIL import Image
from RealESRGAN import RealESRGAN
from flask import Flask, request, jsonify, send_file
import io
import logging

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

app = Flask(__name__)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logger.info(f'Using device: {device}')

model2 = RealESRGAN(device, scale=2)
model2.load_weights('weights/RealESRGAN_x2.pth', download=True)
logger.info('Model x2 loaded successfully')

model4 = RealESRGAN(device, scale=4)
model4.load_weights('weights/RealESRGAN_x4.pth', download=True)
logger.info('Model x4 loaded successfully')

model8 = RealESRGAN(device, scale=8)
model8.load_weights('weights/RealESRGAN_x8.pth', download=True)
logger.info('Model x8 loaded successfully')

def inference(image, size):
    global model2, model4, model8

    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        logger.info('CUDA cache cleared')

    logger.info(f'Starting inference with scale {size}')
    
    try:
        if size == '2x':
            result = model2.predict(image.convert('RGB'))
        elif size == '4x':
            result = model4.predict(image.convert('RGB'))
        else:
            width, height = image.size
            if width >= 5000 or height >= 5000:
                return None, "The image is too large."
            result = model8.predict(image.convert('RGB'))
        logger.info(f'Inference completed for scale {size}')
    except torch.cuda.OutOfMemoryError as e:
        logger.error(f'OutOfMemoryError: {e}')
        logger.info(f'Reloading model for scale {size}')
        
        if size == '2x':
            model2 = RealESRGAN(device, scale=2)
            model2.load_weights('weights/RealESRGAN_x2.pth', download=False)
            result = model2.predict(image.convert('RGB'))
        elif size == '4x':
            model4 = RealESRGAN(device, scale=4)
            model4.load_weights('weights/RealESRGAN_x4.pth', download=False)
            result = model4.predict(image.convert('RGB'))
        else:
            model8 = RealESRGAN(device, scale=8)
            model8.load_weights('weights/RealESRGAN_x8.pth', download=False)
            result = model8.predict(image.convert('RGB'))
        logger.info(f'Model reloaded and inference completed for scale {size}')

    return result, None

@app.route('/upscale', methods=['POST'])
def upscale():
    if 'image' not in request.files:
        logger.warning('No image uploaded')
        return jsonify({"error": "No image uploaded"}), 400
    
    image_file = request.files['image']
    size = request.form.get('size', '2x')

    try:
        image = Image.open(image_file)
        logger.info(f'Image uploaded and opened successfully')
    except Exception as e:
        logger.error(f'Invalid image file: {e}')
        return jsonify({"error": "Invalid image file"}), 400

    result, error = inference(image, size)
    
    if error:
        logger.error(f'Error during inference: {error}')
        return jsonify({"error": error}), 400
    
    img_io = io.BytesIO()
    result.save(img_io, 'PNG')
    img_io.seek(0)
    logger.info('Image processing completed and ready to be sent back')

    return send_file(img_io, mimetype='image/png')

if __name__ == '__main__':
    logger.info('Starting the Flask server...')
    app.run(host='0.0.0.0', port=5000)