#!/usr/bin/env python

import gunicorn.app.base
import sys
from predictor import Predictor, Prediction

class CustomUnicornApp(gunicorn.app.base.BaseApplication):
    """ 
    This gunicorn app class provides create and exit callbacks for workers,     
    and runs gunicorn with a single worker and multiple gthreads
    """
    def __init__(self, create_app_callback, exit_app_callback, host_port):
        self._configBind = host_port
        self._createAppCallback = create_app_callback
        self._exitAppCallback = exit_app_callback
        super().__init__()

    @staticmethod
    def exitWorker(arbiter, worker):
        # worker.app provides us with a reference to "self", and we can call the 
        # exit callback with the object created by the createAppCallback:
        self = worker.app
        self._exitAppCallback(self._createdApp)

    def load_config(self):
        self.cfg.set("bind", self._configBind)
        self.cfg.set("worker_class", "gthread")
        self.cfg.set("workers", 1)
        self.cfg.set("threads", 4)
        self.cfg.set("worker_exit", CustomUnicornApp.exitWorker)
        # Try to uncomment and make 10 requests, to test correct restart of worker:
        # self.cfg.set("max_requests", 10) 

    def load(self):
        # This function is invoked when a worker is booted
        self._createdApp = self._createAppCallback()
        return self._createdApp
    

# --- index.html contents

html = """<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>Cell Growth Classifier</title>
<script>
function uploadImage() {
    var formData = new FormData();
    var imageFile = document.getElementById("imageInput").files[0];
    formData.append("image", imageFile);

    fetch('/upload', {
        method: 'POST',
        body: formData
    })
    .then(response => response.json())
    .then(data => {
        document.getElementById("results").textContent = `${data.results}`;
        var image = document.createElement("img");
        image.src = URL.createObjectURL(imageFile);
        image.onload = function() {
            URL.revokeObjectURL(image.src) // Free up memory
        }
        document.getElementById("imageDisplay").innerHTML = '';
        document.getElementById("imageDisplay").appendChild(image);
    })
    .catch(error => {
        console.error('Error:', error);
        document.getElementById("results").textContent = "An error occurred while uploading the image.";
    });
}
</script>
</head>
<body>

<h2>Cell growth classifier demo</h2>
<input type="file" id="imageInput" accept="image/*"> <br><br>
<button onclick="uploadImage()">Upload</button>

<h3>Image:</h3>
<div id="imageDisplay"></div>

<h3>Results:</h3>
<div id="results"></div>

</body>
</html>
"""

# ---- 

from PIL import Image
from bottle import Bottle, request, response
import threading
import io

def startServer(index_html:str, predictor:Predictor, host_port:str):
    def create():
        app = Bottle()
        lock = threading.Lock()
    
        @app.route("/")
        def getIndex():
            # Serve static content, no lock protection necessary:
            return index_html

        @app.route('/upload', method='POST')
        def upload_image():
            upload = request.files.get('image')
            if upload is None:
                return "No image uploaded"
            
            # Read the image file in bytes and open it with Pillow
            image_bytes = io.BytesIO(upload.file.read())
            try:
                with Image.open(image_bytes) as img:
                    pred:Prediction = app.predictor.predict(img)
                    return {"results": str(pred)}
            except IOError:
                response.status = 400
                return "Invalid image file"
  
        # Store predictor in the bottle app object
        app.predictor = predictor
        return app

    def exit(app):
        # Get the service through the app object and save state
        pass

    CustomUnicornApp(create, exit, host_port).run()

def usage():
    print("""Usage:
    server.py modelfile host_port

Example:      
    ./server.py cells_2.pth localhost:8001
""")

if __name__ == "__main__":
    if len(sys.argv) != 3:
        usage()
    else:
        p = Predictor(sys.argv[1])
        startServer(html, p, sys.argv[2])

