import frontmatter
import gradio as gr
import json
import spaces

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

model_id = "AverageBusinessUser/aidapal"
filename = "aidapal-8k.Q4_K_M.gguf"

print("Downloading model")

tokenizer = AutoTokenizer.from_pretrained(model_id, gguf_file=filename)
model = AutoModelForCausalLM.from_pretrained(
    model_id, gguf_file=filename, device_map="auto"
)

example = """int __fastcall sub_B0D04(int a1, int a2)
{
  unsigned int v2; // r4
  int result; // r0

  v2 = a1 + a2;
  if ( __CFADD__(a1, a2) )
    return 0;
  result = _libc_alloca_cutoff();
  if ( v2 <= 0x1000 )
    return result | 1;
  return result;
}"""

examples = [json.loads(line)["input"] for line in open("gpt4_juiced_dataset.json", "r")]

# Then create the pipeline with the model and tokenizer
pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer)

# TEMPLATE """{{ .System }}
# [INST]
# {{ .Prompt }}
# [/INST]
# """
# SYSTEM """<s>[INST]You are an expert at analyzing code that has been decompiled with IDA Hex Rays into IDA Hex Rays pseudocode. As a IDA Hex Rays pseudocode analyzer, you will be provided code that may or may not have symbols and variable names. You will analyze the IDA Hex Rays pseudocode and explain exactly what each line is doing. Then you will review your analysis and determine potential name for the function and variables within the function. Your task is use your knowledge of reverse engineering, IDA Hex Rays pseudocode, and C to assist the user with analysis and reverse engineering. Provide a detailed description of the Hex Rays pseudocode to the user explaining what the code does, suggest a function name based on the analysis of the pseudocode, and new variable names based on the analysis of the code. Only respond with valid JSON using the keys 'function_name','comment', and an array 'variables'. Values should use plain ascii with no special characters.
# Analyze the following IDA Hex Rays pseudocode and generate a valid JSON object containing the keys 'function_name','comment', and an array 'variables' explaining what the code does, suggest a function name based on the analysis of the code, and new variable names based on the analysis of the code.[/INST]</s>
# """

system = """<s>[INST]You are an expert at analyzing code that has been decompiled with IDA Hex Rays into IDA Hex Rays pseudocode. As a IDA Hex Rays pseudocode analyzer, you will be provided code that may or may not have symbols and variable names. You will analyze the IDA Hex Rays pseudocode and explain exactly what each line is doing. Then you will review your analysis and determine potential name for the function and variables within the function. Your task is use your knowledge of reverse engineering, IDA Hex Rays pseudocode, and C to assist the user with analysis and reverse engineering. Provide a detailed description of the Hex Rays pseudocode to the user explaining what the code does, suggest a function name based on the analysis of the pseudocode, and new variable names based on the analysis of the code. Only respond with valid JSON using the keys 'function_name','comment', and an array 'variables'. Values should use plain ascii with no special characters.
Analyze the following IDA Hex Rays pseudocode and generate a valid JSON object containing the keys 'function_name','comment', and an array 'variables' explaining what the code does, suggest a function name based on the analysis of the code, and new variable names based on the analysis of the code.[/INST]</s>
"""


@spaces.GPU
def predict(code):
    prompt = f"""{system}
[INST]
{code}
[/INST]
"""

    print(f"Prompt: {repr(prompt)}")
    print(f"Tokenized: {tokenizer.tokenize(prompt)}")
    pipe_out = pipe(
        prompt,
        do_sample=True,
        top_k=100,
        top_p=0.09,
        temperature=1.2,
        repetition_penalty=1.1,
        return_full_text=False,
        max_length=4096,
    )
    print(f"Pipe out: {repr(pipe_out)}")

    raw_output = pipe_out[0]["generated_text"]
    output = raw_output
    if output.startswith("```json\n"):
        output = output[8:]

    json_output = json.dumps([])
    try:
        json.loads(output)
        json_output = output
    except Exception:
        pass

    print(f"JSON output: {repr(json_output)}")

    return json_output, raw_output


demo = gr.Interface(
    fn=predict,
    inputs=gr.Text(label="Hex-Rays decompiler output"),
    outputs=[gr.JSON(label="Aidapal Output as JSON"), gr.Text(label="Raw Aidapal Output")],
    description=frontmatter.load("README.md").content,
    examples=examples
)
demo.launch()