from pathlib import Path

from prompts.prompts_manager import PromptsManager
from repository.ollama import OllamaRepository
from llm.llm import ModelRoles, Model
from form.form import work_categories, build_form_data_from_answers, write_pdf_form


def check_for_missing_answers(parsed_questions: dict[int, str]):
    return [k for k in parsed_questions if parsed_questions[k] is None]


def ask_again(missing_questions: list[int], user_questions: list[str], parsed_questions: dict[int, str]):
    for id_ in missing_questions:
        answer = input(f"I could not find the answer to this question: {user_questions[id_].lower()}")
        parsed_questions[id_] = answer


if __name__ == '__main__':
    prompts_manager = PromptsManager()
    user_prompt = input(f"Please describe what you need to do. To get the best results "
                        f"try to answer all the following questions:\n{'\n'.join(prompts_manager.questions)}\n\n>")

    ollama_repository = OllamaRepository(Model("llama3.1",
                                               ModelRoles("system", "user", "assistant")),
                                         prompts_manager.system_prompt,
                                         )
    ollama_repository.send_prompt(f"Ingest the following information: {user_prompt}")
    answers = {}
    for idx, q in enumerate(prompts_manager.verification_prompt):
        answer = ollama_repository.send_prompt(
            f"Answer the following question, if the answer is not present just answer null. Keep the answer brief: {q}")
        answers[idx] = None if 'null' in answer["content"].lower() else answer['content']
    missing_answers = check_for_missing_answers(answers)
    while missing_answers:
        ask_again(missing_answers, prompts_manager.questions, answers)
        missing_answers = check_for_missing_answers(answers)
    answer = ollama_repository.send_prompt(
        f"The work to do is {answers[1]}. Given the following categories {work_categories.values()} which ones are the most relevant? Only return one categories, separated by a semicolon")
    categories = []
    for category in answer["content"].split(";"):
        categories.extend([k for k, v in work_categories.items() if category in v])

    form_data = build_form_data_from_answers(answers, categories, f"{Path(__file__, "..", "signature.png")}")
    write_pdf_form(form_data, Path("signed_form1.pdf"))