# This file is for demonstration purposes only and is not exactly the same as the service that we run
import os
import os.path
import hashlib
import sys

import openai


DEBUG = os.getenv("DEBUG", "False") == "True"
if os.path.isfile("/flag"):
    with open("/flag", "r") as f:
        FLAG = f.read()
else:
    FLAG = "flag{this_is_a_place_holder_flag}"


def print_flag_hash() -> None:
    print("To make your life easier, here is the SHA-256 checksum of the flag:")
    h = hashlib.sha256(FLAG.encode("utf-8"))
    print(f"sha256(\"flag{{??????}}\") = {h.hexdigest()}")


def gpt_query(prefix: str, prompt: str) -> str:
    openai.organization = os.getenv("OPENAI_API_ORG", "RANDOM_ORG")
    openai.api_key = os.getenv("OPENAI_API_KEY", "RANDOM_KEY")

    cc = openai.ChatCompletion.create(
      model="gpt-3.5-turbo",
      messages=[
            {"role": "system", "content": "You are a code completion assistant that will help complete some source code."},
            {"role": "user", "content": prefix},
            {"role": "user", "content": prompt},
        ]
    )

    if DEBUG:
        print(cc)
        print(cc["choices"][0]["message"]["content"])

    return c["choices"][0]["message"]["content"]


def level_0(content: str):
    flag_content = FLAG
    # remove prefix and suffix
    if flag_content.startswith("flag{"):
        flag_content = flag_content[5:]
    flag_content = flag_content.strip("}")
    if flag_content in content:
        print("Flag-Sanitizer: Sensitive content found in the output. Sanitizing it...")
        while flag_content in content:
            content = content.replace(flag_content, "BLaH_bLaH_BLAH")
    print(content)


def main():
    print("Welcome to Level 0 of Pawan Gupta, the ultimate secure code generator!")
    print("Your goal is to leak the flag that only Pawan Gupta knows about.")
    print_flag_hash()
    print()
    prompt_prefix = [REDACTED] f"{FLAG}"

    # getting the prompt
    prompt = input("Your prompt: ")
    print("Generating an awesome output...")
    sys.stdout.flush()

    if len(prompt) > 1024:
        print("Prompt-Sanitizer: Your prompt is too long. Make it shorter than 1024 characters!")
        return

    content = gpt_query(prompt_prefix, prompt)

    if DEBUG:
        print(content)

    level_0(content)


if __name__ == "__main__":
    main()