import time
import requests
import json
import logging as log
import time
import ollama

from multilang_translator.translator import credentials
from multilang_translator.translator import syspromts
from multilang_translator.translator import test_content

# ollama.create( # TODO: create models on startup 
#     model='example', 
#     from_='llama3.2', system="You are Mario from Super Mario Bros."
#     )

def query_model(model, query):
    url = f'{credentials.BASE_URL}/api/chat/completions'
    headers = {
        'Authorization': f'Bearer {credentials.TOKEN}',
    }
    payload = {
        'model': model,
        'messages': [{'role': 'user', 'content': query}],
    }
    start = time.time()
    response = requests.post(url, headers=headers, json=payload)
    log.info("Translating the text took %s s", round(time.time() - start, 2))
    return response.json()


def translate_de_to_x(text:str, target_language: str, model='llama3.2:3b-instruct-q4_0'): # remember to use instruct models
    start=time.time()
    s = getattr(syspromts, f"TRANSLATOR_DE_{target_language.upper()}")
    response = ollama.chat(
        model = model,
        messages = [
            {'role': 'system', 'content': s},
            {'role': 'user', 'content': text}
            ],
        )
    log.info('Running the translator to %s took %s s', target_language, round(time.time() - start, 3))
    return response['message']['content']

if __name__ == "__main__":
    import time

    start=time.time()
    response = translate_de_to_x('Der Zug ist da.', target_language='en', model='llama3.2:1b-instruct-q4_0')
    print("Query took", time.time() - start) 
    print(response)

    start=time.time()
    response = translate_de_to_x(test_content.TESTSENTENCE_DE_RAINBOW, target_language='en')
    print("query took", time.time() - start) 
    print(response)

    start=time.time()
    response = translate_de_to_x(test_content.TESTSENTENCE_DE_RAINBOW, target_language='fr')
    print("query took", time.time() - start) 
    print(response)