#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-or-later

import json
import graphviz
import sys
from html.parser import HTMLParser

EDGE_LABELS = False
if len(sys.argv) > 1 and sys.argv[1] == '-e':
    EDGE_LABELS = True

with open('quiz1.json', 'r') as f:
    quiz = json.load(f)

with open('quiz2.json', 'r') as f:
    quiz |= json.load(f)

with open('scorecard.json', 'r') as f:
    scorecard = json.load(f)

dot = graphviz.Digraph('catgirl-or-puppygirl')
dot.attr('node', style='filled', fillcolor='white')


class QuizHTMLParser(HTMLParser):
    def __init__(self):
        self.question = None
        self.answers = []
        self.answer_index = None
        self.next_data_is_question_text = False
        self.next_data_is_answer_text = False
        HTMLParser.__init__(self)

    def handle_starttag(self, tag, attrs):
        match tag:
            case 'div':
                for attr in attrs:
                    if attr == ('class', 'question_text'):
                        self.next_data_is_question_text = True
            case 'label':
                for attr in attrs:
                    if attr == ('class', 'answer_row'):
                        self.answers.append({})
                        if self.answer_index is None:
                            self.answer_index = 0
                        else:
                            self.answer_index += 1
            case 'input':
                if self.answer_index is None:
                    return
                if ('name', 'SelectedAnswers[]') not in attrs:
                    return
                for k, v in attrs:
                    if k == 'value':
                        self.answers[self.answer_index]['id'] = v
            case 'span':
                if self.answer_index is None:
                    return
                for attr in attrs:
                    if attr == ('class', 'answer_text'):
                        self.next_data_is_answer_text = True

    def handle_data(self, data):
        if self.next_data_is_answer_text:
            self.answers[self.answer_index]['text'] = data
            self.next_data_is_answer_text = False
        elif self.next_data_is_question_text:
            self.question = data
            self.next_data_is_question_text = False


questions = []
all_answers = []
for v in quiz.values():
    html = v['QuestionHtml']
    parser = QuizHTMLParser()
    parser.feed(html)
    all_answers.append(parser.answers)
    questions.append(parser.question)

question_colours = [
    'lightgray', 'lightgreen', 'lightblue', 'yellow', 'pink', 'orange'
]

question_colours_dark = [
    'gray', 'green', 'blue', 'goldenrod', 'purple', 'darkorange'
]

i = 0
for question in questions:
    dot.node(f'question{i}', question, shape='box',
             fillcolor=question_colours[i % len(question_colours)])
    i += 1

personality_type_ids = {
    '10189839': 'Dominant catgirl',
    '10189841': 'Submissive catgirl',
    '10189845': 'Submissive puppygirl',
    '10189842': 'Dominant puppygirl',
    '10189846': 'Infertile warthog',
}

for type_id, name in personality_type_ids.items():
    dot.node(type_id, name)

i = 0
for question in scorecard['ScoreCardQuestions']:
    answers = question['PersonalityAnswers']
    if answers is None:
        continue
    for answer in answers:
        for possible_answer in all_answers[i]:
            if int(possible_answer['id']) == answer['PossibleAnswerId']:
                if EDGE_LABELS:
                    for personality_type_id in answer['PersonalityTypeIds']:
                        dot.edge(
                            f'question{i}', str(personality_type_id),
                            label=possible_answer['text'],
                            color=question_colours_dark[i % len(question_colours_dark)],
                            fontcolor=question_colours_dark[i % len(question_colours_dark)],
                            labeldistance='16',
                            decorate='true'
                        )
                else:
                    dot.node(
                        possible_answer['id'], possible_answer['text'],
                        fillcolor=question_colours[i % len(question_colours)]
                    )
                    dot.edge(f'question{i}', possible_answer['id'])
                    for personality_type_id in answer['PersonalityTypeIds']:
                        dot.edge(
                            possible_answer['id'],
                            str(personality_type_id)
                        )
    i += 1

print(dot.source)