#!/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 uploadImages() {
    var formData = new FormData();
    var imageFiles = document.getElementById("imageInput").files;
    var imageFilesArray = [];
    for (var i = 0; i < imageFiles.length; i++) {
        formData.append("images", imageFiles[i]);
        imageFilesArray.push({name: imageFiles[i].name, file: imageFiles[i]});
    }

    fetch('/upload', {
        method: 'POST',
        body: formData
    })
    .then(response => response.json())
    .then(data => {
        var results = data.results;
        var sortedKeys = Object.keys(results).sort(); // Sort the filenames
        var table = document.getElementById("resultsTable");
        table.innerHTML = ""; // Clear the table first

        // Create table header
        var header = table.createTHead();
        var headerRow = header.insertRow(0);
        headerRow.insertCell(0).textContent = "Image";
        headerRow.insertCell(1).textContent = "Filename";
        headerRow.insertCell(2).textContent = "Analysis";

        // Create table body
        var tbody = table.createTBody();

        // Insert a row in the table for each sorted image result
        sortedKeys.forEach((imageName) => {
            var imageFile = imageFilesArray.find(file => file.name === imageName);
            var row = tbody.insertRow(-1);
            var cellImage = row.insertCell(0);
            var image = document.createElement("img");

            if (imageFile) {
                image.src = URL.createObjectURL(imageFile.file);
                image.onload = function() {
                    URL.revokeObjectURL(this.src) // Free up memory
                }
            } else {
                image.alt = "Image not found";
            }
            image.style.width = '100px'; // Set the image size
            cellImage.appendChild(image);

            row.insertCell(1).textContent = imageName;
            row.insertCell(2).textContent = results[imageName].text;
        });

        // Set the model used
        document.getElementById("modelUsed").textContent = "Model used: " + data.model;
    })
    .catch(error => {
        console.error('Error:', error);
        document.getElementById("results").textContent = "An error occurred while uploading the images.";
    });
}
</script>
<style>
  table {
    border-collapse: collapse;
    margin-top: 20px;
  }
  thead, td {
    border: 1px solid #ddd;
    padding: 8px;
    text-align: left;
  }
  thead {
    background-color: #f2f2f2;
    color: #333;
    font-weight: bold;
  }
  td {
    font-size: 0.9em;
  }
  img {
    width: 100px; /* Adjust the size of the image if necessary */
    height: auto;
  }
  #modelUsed {
    margin-top: 20px;
  }
</style>

</head>
<body>

<h2>Cell Growth Classifier</h2>
<input type="file" id="imageInput" accept="image/*" multiple>
<button onclick="uploadImages()">Analyze</button>

<p id="modelUsed"></p>
<table id="resultsTable">
    
</table>

</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():
            uploads = request.files.getall('images')
            results = {}
            for upload in uploads:
                # 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)
                        results[upload.filename] = pred.getDict()
                except IOError:
                    response.status = 400
                    return "Invalid image file"
            return {"model": app.predictor.modelName, "results": results}
  
        # 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])

