from form.form import work_categories
from prompts.prompts_manager import Questions


class LlmParser:

    @classmethod
    def parse_verification_prompt_answers(cls, llm_answer) -> dict[Questions, str | None]:
        print(f"llm answer: {llm_answer}")
        answers = {}
        i = 0
        question_id = 0
        lines = [l for l in llm_answer.split("\n") if len(l.strip()) > 0 and not l.strip().endswith("?")]
        while i < len(lines):
            line = lines[i].strip()
            if len(line) == 0:
                i += 1
            elif "null" in lines[i]:
                answers[Questions(question_id)] = None
                i += 1
                question_id += 1
            elif ":" in lines[i]:
                answers[Questions(question_id)] = line.split(":")[1]
                i += 1
                question_id += 1
            else:
                answers[Questions(question_id)] = line
                i += 1
                question_id += 1
        return answers

    @classmethod
    def parse_get_categories_answer(cls, category_answer) -> list[str]:
        categories = []
        for category in category_answer.split(";"):
            categories.extend([k for k, v in work_categories.items() if category.lower() in v.lower()])
        return categories