import json
from json import JSONDecodeError

from repository.ollama import OllamaRepository
from schema import ModelRoles

if __name__ == '__main__':
    with open("questions.txt") as questions_file:
        questions = questions_file.read()
    with open("system_prompt.txt") as system_prompt_file:
        system_prompt = system_prompt_file.read()
    with open("verification_prompt.txt") as verification_prompt_file:
        verification_prompt = verification_prompt_file.read()
    verification_prompt = verification_prompt.replace("{questions}", questions)
    user_prompt = input(f"Please describe what you need to do. To get the best results "
          f"try to answer the following questions:\n{questions}\n\n>")

    ollama_repository = OllamaRepository("llama3.1", system_prompt,
                                         ModelRoles("system", "user", "assistant"))

    ollama_repository.send_prompt(f"Ingest the following information: {user_prompt}")
    answer = ollama_repository.send_prompt(verification_prompt)
    json_data_start = answer["content"].index("{")
    json_data_stop = answer["content"].rindex("}")

    answers = "[" + answer["content"][json_data_start:json_data_stop+1] + "]"
    try:
        json_answers = json.loads(answers)
        print(json_answers)
    except JSONDecodeError as e:
        print(f"unable to parse \n {answers}")