from flask import Flask, render_template, request, make_response, redirect
from freecaptcha import captcha
import uuid
import os
from cryptography.fernet import Fernet
import jwt
import base64

# Generate a new AES key
key = Fernet.generate_key()

JWT_SECRET_KEY = 'obviously this needs to be an env variable secret thingie.'

# Create an instance of the Fernet cipher with the key
cipher = Fernet(key)

# Encrypt a message
message = b'This is a secret message.'
ciphertext = cipher.encrypt(message)

# Decrypt the ciphertext
app = Flask(__name__, static_url_path='', static_folder='static',)


@app.route("/captcha", methods=['GET', 'POST'])
def captcha_handler():
    if request.method == "POST":
        captcha_attempt = len(list(request.form))
        token = request.cookies.get('Authorization').split('Bearer ')[-1]
        try:
            # TODO: set JWT to expire very soon.
            payload = jwt.decode(token, JWT_SECRET_KEY, algorithms=['HS256'])
            b64_and_encrypted_correct_answer = payload['encrypted_correct_answer']
            n = payload['n']
            encrypted_correct_answer_bytes = base64.b64decode(b64_and_encrypted_correct_answer)
            correct_answer = cipher.decrypt(encrypted_correct_answer_bytes).decode('utf-8').split('|||')[0]
            ## Redirect to the original page the user wanted - with a token letting that they can validate from us that says that the user passed a specific captcha attempt (we will sign the attempt with a code we give them with the captcha, like an id, so they know it was that specific attempt)
            return f'''
            The correct answer was {correct_answer}
            You flipped it {int(captcha_attempt) % n}
            '''

        except jwt.ExpiredSignatureError:
            return 'Token expired. Please log in again.'

        except jwt.InvalidTokenError:
            return 'Invalid token. Authentication required.'
        ## Get their guess from POST body
        ## Get JWT cookie
        ## Decrypt payload
        ## Check if guess (from post body) == correct answer (from decrypted JWT payload)
        ## If so: give them a new JWT for winning the CAPTCHA, and forward them to the URL they originally wanted to access precaptcha (just like oauth does)
        ## If not: Redirect them to the GET version of this same URL, with warning enabled to tell them they failed
    if request.method == "GET":
        image_path = captcha.random_image()
        n = 6
        answer, options = captcha.captchafy(image_path, n)
        print('the correct answer is: ', answer)
        # remember to store the salt since we'll need it when we compare the hashes
        salt = uuid.uuid4()
        plaintext_bytes = bytes(str(answer) + '|||' + str(salt), 'utf-8')
        encrypted_bytes = cipher.encrypt(plaintext_bytes)

        ciphertext = base64.b64encode(encrypted_bytes).decode('utf-8')
        token = jwt.encode({
            'encrypted_correct_answer': ciphertext,
            'salt': str(salt),
            'n': n
        }, JWT_SECRET_KEY, algorithm='HS256')

        # Set the Authorization header cookie with the JWT
        resp = make_response(captcha.generate_captcha_html(list(options)))
        resp.set_cookie('Authorization', f'Bearer {token}')

        return resp
        # Set JWT with encrypted answer as HTTP cookie
        # issue: they could compare our encrypted answer with how we encrypt all numbers 1-6, since they will have seem them before.
        # We could solve this with a salt, but then we have to store salts. We could include the salt in plaintext, but then they could build a rainbow table of our guess + salts (since there are only 6 possible guesses)
        # the solution is to use salts big enough that we don't have to fear them repeating / being turned into a rainbow table.
        # We will use UUID's as the salts.
        #
        # Anyway, we pass the data to our Jinja template and render it.
    else:
        return "Unsupported HTTP method."
    # Flask should take care of unsupported methods for us.