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']
            encrypted_correct_answer_bytes = base64.b64decode(b64_and_encrypted_correct_answer)
            correct_answer = cipher.decrypt(encrypted_correct_answer_bytes).decode('utf-8').split('|||')[0]
            return f'''
            The correct answer was {correct_answer}
            You flipped it {captcha_attempt}
            '''

        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()
        answer, options = captcha.captchafy(image_path)
        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)
        }, 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.
    # Flask should take care of unsupported methods for us.



    ## Handle cookie

    ## Get random image
    ## Generate 
    return captcha.generate_captcha_html(os.listdir('static/images/'))

@app.route("/captcha_old", methods=['GET', 'POST'])
def login():
    # This means they just submitted a CAPTCHA
    # We need to see if they got it right
    incorrect_captcha = False
    if request.method == 'POST':
        captcha_guess = len(list(request.form))
        print(request.form.get('captcha'))
        # What if they POST with the cookie below absent? Uh oh...
        captcha_cookie = request.cookies.get('freecaptcha_cookie')

        real_answer = captcha_solutions.get(captcha_cookie, None)
        if real_answer is not None:
            if captcha_guess == int(real_answer):
                captcha_solved.append(captcha_cookie)
                return redirect("/", code=302)
            else:
                incorrect_captcha = True
                

    # Select an image
    image_path = captcha.random_image()

    # Generate list of rotated versions of image
    # and save which one is correct
    # change answer to be the number of turns needed?
    answer, options = captcha.captchafy(image_path)
    print(answer)

    # Provide the CAPTCHA options to the web page using the CAPTCHA
    resp = make_response(render_template("index.html", captcha_options=options, incorrect_captcha=incorrect_captcha))
    resp = make_response(captcha.generate_captcha_html(list(options)))
    # Track this user with a cookie and store the correct answer
    # by linking the cookie with the answer, we can check their answer later
    freecaptcha_cookie = str(uuid.uuid4())
    resp.set_cookie('freecaptcha_cookie', freecaptcha_cookie)
    captcha_solutions[freecaptcha_cookie] = answer

    return resp