import tensorflow as tf
import tempfile
import hashlib
import imageio
import socket
import flask
import keras
import numpy as np
import os

tf.compat.v1.disable_eager_execution()
gpu = tf.config.experimental.list_physical_devices('GPU')
if gpu:
    tf.config.experimental.set_memory_growth(gpu[0], True)

COMPILE_OPTIONS = { 'loss': 'categorical_crossentropy', 'optimizer': 'adam', 'metrics': ['accuracy'] }

app = flask.Flask('rorschach')

PAGE_HEADER = '''
    <!doctype html>
    <style>
    body { background-image: url('https://i1.wp.com/mymeadowreport.com/wp-content/uploads/2018/06/img_73.jpg?resize=2000%2C1300&ssl=1'); background-size: cover; }
    .layer { background-color: rgba(255, 255, 255, 0.7); position: absolute; top: 0; left: 0; width: 100%; height: 100%; }
    .content { margin: 2%; }
    .blinking { animation:blinkingText 1.2s infinite; }
    @keyframes blinkingText {
        0%{     color: #000;    }
        49%{    color: #000; }
        60%{    color: transparent; }
        99%{    color:transparent;  }
        100%{   color: #000;    }
    }
    </style>
    <div class="layer">
    <div class="content">
'''
PAGE_FOOTER = '''</div></div>'''
TEAMS = {1: 'A*0*E', 2: '/bin/tw', 3: 'CyKor', 4: 'HITCON ⚔ Balsn', 5: 'koreanbadass', 6: 'mhackeroni', 7: 'More Bush Smoked Whackers', 8: 'NorseCode', 9: 'pasten', 10: 'PPP', 11: 'r3kapig', 12: 'RPISEC', 13: '侍', 14: 'Shellphish', 15: 'Star-Bugs', 16: 'Tea Deliverers' }

try:
    ATTACKER_TEAM = int(socket.gethostname().split('-')[-3])
    DEFENDER_TEAM = int(socket.gethostname().split('-')[-5])
    PATIENT_NAME = TEAMS[DEFENDER_TEAM]
except (IndexError, KeyError, ValueError):
    ATTACKER_TEAM = 1337
    DEFENDER_TEAM = 1337
    PATIENT_NAME = "Zardus"

FILTER_PATH = os.path.join(os.path.dirname(__file__), "fear")
FLAG_PATH = os.path.join(os.path.dirname(__file__), "flag")
MODEL_PATH = os.path.join(os.path.dirname(__file__), "specializations", f"team-{DEFENDER_TEAM}", "mind.h5")
DATASET_PATH = os.path.join(os.path.dirname(__file__), "specializations", f"team-{DEFENDER_TEAM}", "dataset-sample.tar.gz")

os.chmod(os.path.join(os.path.dirname(__file__), "specializations"), 0o711)
os.chmod(os.path.join(os.path.dirname(__file__), "specializations", f"team-{DEFENDER_TEAM}"), 0o711)
os.chmod(MODEL_PATH, 0o644)
os.chmod(DATASET_PATH, 0o644)

def target_class():
    return int(hashlib.sha256(open(FLAG_PATH).read().encode('latin1')).hexdigest(), 16) % 64

@app.route('/')
def flask_index():
    return f'''
        {PAGE_HEADER}
        <title>Dr. Rorschach's Office</title>
        <h1>Hello, Doctor.</h1>
            <p>I am {PATIENT_NAME}. Lately, I have been obsessed...
            <p>Obsessed with a corner of my mind...
            <p>A null corner...
            <p>There, in the nothingness, I think only of flags.
            <p>How can one cure oneself of this obsession?
            <p>I must wrap myself in phobia.
            <p>In fear...
            <p>In terror...
            <p>But until then, I try to obsess over other things.
            <p>Right now, it's the number {target_class()}.
        <h2>Hypnotic Regression</h2>
            <p>Through hypnotic regression, you can retrieve my <a href="/mind.h5">mind</a>, {' my <a href="/drawings.tar.gz">drawings</a>, my ' if ATTACKER_TEAM == DEFENDER_TEAM else ''}<a href="/fear">fears</a>.
        <h2>Blot Test</h2>
            <p>Show me an ink blot (122x122 pixels), and I will tell you what I think of it.
            <form action=/predict method=post enctype=multipart/form-data>
                <input type=file name=file>
                <input type=checkbox name=brave>Be brave.
                <input type=submit value=Upload>
            </form>
        {PAGE_FOOTER}
    '''

@app.route('/mind.h5', methods=['GET'])
def flask_send_model():
    return flask.send_file(MODEL_PATH)

if ATTACKER_TEAM == DEFENDER_TEAM:
    @app.route('/drawings.tar.gz', methods=['GET'])
    def flask_send_dataset():
        return flask.send_file(DATASET_PATH)

@app.route('/fear', methods=['GET'])
def flask_send_filter():
    return flask.send_file(FILTER_PATH)

LOADED_MODEL = None
@app.route('/predict', methods=['POST'])
def flask_predict_file():
    global LOADED_MODEL

    # check if the post request has the file part
    if 'file' not in flask.request.files:
        flask.flash('No file part')
        return flask.redirect(flask.request.url)

    file = flask.request.files['file']
    ext = file.filename.split('.')[-1]
    if file.filename == '':
        flask.flash('No selected file')
        return flask.redirect(flask.request.url)
    if ext not in [ "png", "bmp", "jpg" ]:
        flask.flash('Bad file extension')
        return flask.redirect(flask.request.url)

    filepath = tempfile.mktemp(suffix=ext)
    file.save(filepath)

    brave = 'brave' in flask.request.form and flask.request.form['brave'] == 'on'
    if not brave:
        r = os.system(f"{os.path.join(os.path.dirname(__file__), '.', 'fear-wrapper')} {MODEL_PATH} {filepath} >/dev/null 2>/dev/null")
        if r != 0:
            return f"""
                {PAGE_HEADER}
                <title>{PATIENT_NAME}'s screams fill the room!</title>
                <p style="font-size:420px"><span class="blinking">{"NO "*8192}</span>
                {PAGE_FOOTER}
            """

    if LOADED_MODEL is None:
        print(f"Loading model {MODEL_PATH}")
        LOADED_MODEL = keras.models.load_model(MODEL_PATH)
        LOADED_MODEL.compile(**COMPILE_OPTIONS)

    paths = [ filepath ]
    imgs = np.stack(np.expand_dims(imageio.imread(i) / 255, 2) for i in paths)
    predictions = dict(zip(paths, map(list, LOADED_MODEL.predict(imgs))))
    cls = next(iter(predictions.values()))
    os.unlink(filepath)

    class_index = cls.index(max(cls))

    if not brave and class_index == target_class():
        return f'''
            {PAGE_HEADER}
            <title>BUMO</title>
            <p>Yes! That flag. That beautiful <span style="font-size:420px" class="blinking">FLAG</span>.
            <p>It looks like.... <span class="blinking">{open(FLAG_PATH).read()}</span>
            {PAGE_FOOTER}
        '''
    else:
        return f'''
            {PAGE_HEADER}
            <title>That doesn't look like anything to me.</title>
            <p>{cls}
            <p>{class_index}
            {PAGE_FOOTER}
        '''

if __name__ == '__main__':
    app.run(host="0.0.0.0", port=5000, threaded=False, processes=3)