#!/home/brickman/anaconda3/bin/python
import requests
import json
import requests
from markdown import markdown
from bs4 import BeautifulSoup
from rich.console import Console
from rich.markdown import Markdown
import sys, os
import tempfile


try:
    if len(sys.argv[1]) > 1:
        question = sys.argv[1]
    else:
        print("Error: Please specifice a question")
        exit()
except:
    print("Error: Please specifice a question")
    exit()

def codeparse(text_code):
    html = markdown(text_code, output_format="html5")
    soup = BeautifulSoup(html, "html.parser")
    for code in soup.findAll('code'):
        code = str(code)
        code = code.replace("<code>bash\n","")
        code = code.replace("<code>","")
        code = code.replace("</code>","")
        print(code)


def check_substring(main_string, substring):
    return substring in main_string

def generate_text(model, prompt, system):
    url = "http://localhost:11434/api/generate"
    data = {
        "model": model,
        "prompt": prompt,
        "system": system,
        "stream": False,
        "options": {
            "temperature": 0.6,
        }
    }
    response = requests.post(url, json=data)
    text = json.loads(response.text)
    return text["response"]

out = ""

count = 0

while not check_substring(out, "CODETRUE"):
    #print('Thinking....')
    text = generate_text("codellama:7b", f"You are Linux Assistant. I will ask you a question in natural language and you will return a single bash code block. Make sure that you keep the result as brief as posible.\n\nPrompt:\n{question}", """You are Linux Assistant. I will ask you a question in natural language and you will return a single bash code block. The format should be as follows:
Discription here
```bash
code goes here
```
""")

    #text = 'The `sudo` command takes an optional argument `-h` which stands for "help". Running `sudo -h | more` will display the help message for the `sudo` command, which includes a description of its usage and options.'

    out = generate_text("codellama:7b", f"""Check the following prompt and see if it lines up with the code.
Look at the prompt and then the resulting bash code. If the result is in the correct format as speciffied below
return CODETRUE if the format is wrong or if the code is incorrect return CODEFALSE. Do not write any code or return CODE_TRUE if you write code that fixes the problem.

Correct format:
Discription here
```bash
code goes here
```

Prompt:
{question}

Bash:
{text}""", "You are a code expert.")
    count+=1
    if count > 4:
        print('\n\nError: Could not find an answer to your problem')
        exit()

# Showing llm output
console = Console()
md = Markdown("# LLM Output\n" + text)
console.print(md)


# printing code blocks
html = markdown(text, output_format="html5")
soup = BeautifulSoup(html, "html.parser")
for code in soup.findAll('code'):
    code = str(code)
    code = code.replace("<code>bash\n","")
    code = code.replace("<code>","")
    code = code.replace("</code>","")
    code = code.replace("&amp;", "&")
    if len(code) < 4:
        continue

    print("\n\n------------------------CODE----------------------------")
    console.print(Markdown(f"```bash\n{code}\n```"))
    print("--------------------------------------------------------")

    answer = input("Would you like to run this code [N/y/e/q]:").lower()
    if answer == "y" or answer == "yes":
        print("\n")
        os.system(code)
        print("")

    elif answer == "e" or answer == "edit":
        file_content = code

        t = tempfile.NamedTemporaryFile(delete=False)
        t.write(file_content.encode())
        t.close()
        try:
            editor = os.environ['EDITOR']
        except KeyError:
            editor = 'nano'
        os.system(f"{editor} {t.name}")
        os.system(f"bash {t.name}")
        os.system(f"rm -rf {t.name}")
    elif answer == "q" or answer == "quit":
        exit()