import json
import os
from datetime import datetime
import requests

HOST = 'localhost:5000'
URI = f'http://{HOST}/api/v1/generate'

def log_stats(file_path, new_data):
    # Check if the file exists, and read its content
    if os.path.exists(file_path):
        with open(file_path, 'r') as file:
            try:
                content = json.load(file)
            except json.JSONDecodeError:
                content = []
    else:
        content = []

    # Get the current date and time
    now = datetime.now()
    # Format the date and time as a string
    date_time_string = now.strftime("%Y-%m-%d %H:%M:%S")
    new_data.append(date_time_string)

    # Append the new data to the content
    content.append(new_data)

    # Write the updated content back to the file
    with open(file_path, 'w') as file:
        json.dump(content, file)


def model_api(request):
    response = requests.post(f'http://{HOST}/api/v1/model', json=request)
    return response.json()

# model info
def model_info():
    response = model_api({'action': 'info'})
    return response['result']['model_name']
    
# simple loader
def model_load(model_name):
    return model_api({'action': 'load', 'model_name': model_name})

def get_completion(prompt, max_tokens = 250, temperature = 0.01):
    request = {
        'prompt': prompt,
        'max_new_tokens': max_tokens,
        'preset': 'None',  
        'do_sample': True,
        'temperature': temperature,
        'top_p': 0.1,
        'typical_p': 1,
        'epsilon_cutoff': 0,  # In units of 1e-4
        'eta_cutoff': 0,  # In units of 1e-4
        'tfs': 1,
        'top_a': 0,
        'repetition_penalty': 1.18,
        'top_k': 40,
        'min_length': 0,
        'no_repeat_ngram_size': 0,
        'num_beams': 1,
        'penalty_alpha': 0,
        'length_penalty': 1,
        'early_stopping': False,
        'mirostat_mode': 0,
        'mirostat_tau': 5,
        'mirostat_eta': 0.1,

        'seed': -1,
        'add_bos_token': True,
        'truncation_length': 2048,
        'ban_eos_token': False,
        'skip_special_tokens': True,
        'stopping_strings': []
    }
    print(URI)
    response = requests.post(URI, json=request)

    if response.status_code == 200:
        result = response.json()['results'][0]['text']
        print(result)
        return result
    else:
        return ""
    
if __name__ == "__main__":

    print(get_completion("hello"))