import json
import requests
import os
import sys
import time
from os.path import expanduser

# get a new github copilot token
def get_token():

    valid = False
    token = {}

    # check if a token exists in the file
    if os.path.exists("token.json"):
        with open("token.json", "r") as f:
            token = json.load(f)

    # check if the token is valid
    if "token" in token and "expires_at" in token:
        if token["expires_at"] > time.time():
            valid = True

    if not valid:
        # read token from file
        
        # read file from home directory
        home = expanduser("~")
        with open(f"{home}/.config/github-copilot/hosts.json", "r") as f:
            terms = json.load(f)
            auth_token = terms["github.com"]["oauth_token"]


            headers = {
                'Authorization': f"Bearer {auth_token}",
            }
            response = requests.get('https://api.github.com/copilot_internal/v2/token', headers=headers)
            
            if response.status_code != 200:
                return None

            token = response.json()

            # save the token in a json file with the expiration time
            with open("token.json", "w") as f:
                json.dump(token, f)

    return token["token"]

# get a suggestion using github copilot
def get_suggestion(prompt, token):
    headers = {
        'Content-Type': 'application/json',
        'Openai-Organization': 'copilot-ghost',
        'OpenAI-Intent': 'github-copilot',
        'Authorization': f'Bearer {token}',
    }

    data = {
        'prompt': prompt,
        'stream': True,
        'temperature': 0.3,
        'max_tokens': 50,
        'top_p': 0.5,
        'n': 1,
        'logprobs': 2,
    }

    response = requests.post('https://copilot-proxy.githubusercontent.com/v1/engines/copilot-codex/completions', headers=headers, data=json.dumps(data))
    # get the first completion from the response
    data = response.text

    # split by \n
    data = data.split("\n")

    # filter where start with data: and remove it
    data = [x.replace("data: ", "") for x in data if x.startswith("data: ")]
    
    # filter [Done]
    data = [x for x in data if x != "[DONE]"]

    choices: dict(int, str) = {}
    for x in data:
        # x to json
        x = json.loads(x)
        c = x["choices"][0]
        
        if c["index"] not in choices:
            choices[c["index"]] = c["text"]
        else:
            choices[c["index"]] += c["text"]
    
    # get the first choice else return "Command not found"
    if 0 in choices:
        choice = choices[0]
        
        # remove unwanted characters from front (#, $, \n)
        while choice.startswith("#") or choice.startswith("$") or choice.startswith("\n") or choice.startswith(" "):
            choice = choice[1:]

        # strip \n
        choice = choice.strip()

        return choice.split("\n")[0]

    return "Command not found"

# cursor_position_char = int(sys.argv[1])

# Read the input prompt from stdin.
# buffer = sys.stdin.read()
buffer = 'change python version to 3.10'

# --- 
prompt_prefix = "!/bin/bash\n\n" + buffer + ":\n"

token = get_token()
print(get_suggestion(prompt_prefix, token))