#include <cstdio>
#include <iostream>
#include <string>
#include <vector>
#include <map>
#include <algorithm>
#include <sstream>
#include <format>

#include "config.h"
#include "lexer.h"
#include "result.h"
#include "parser.h"

struct QuestionElement {
	bool isDash;
	bool isGroup;
    std::string content;
};

std::string MultiElementQuestion::ToString() const {
    std::stringstream choiceOut;
    for (const auto& choice : Choices) {
        char opener;
        if (choice.IsCorrect) {
            opener = '+';
        } else {
            opener = '-';
        }
        choiceOut << opener << " " << choice.Answer << "; ";
    }
	return std::format(
        "<Multiple choice>:{} section: {} id: {}\n{}",
        QuestionText,
        Section,
        ID,
        choiceOut.str()
    );
}

// Automaton for validating token transitions
std::map<TokenType, std::vector<TokenType>> automata;

bool contains(const std::vector<TokenType>& vec, TokenType element) {
    return std::find(vec.begin(), vec.end(), element) != vec.end();
}

// Automata for validating the parser state
std::map<TokenType, std::vector<TokenType>> parserAutomata() {
    std::map<TokenType, std::vector<TokenType>> automata;
    automata[TokenType::TextFragment] = {
        TokenType::QuestionEnd,
        TokenType::ElementDashStart,
        TokenType::ElementPlusStart,
        TokenType::MatchGroupEnd,
        TokenType::SectionIdentifierStart,
        TokenType::SectionStart,
        TokenType::EndOfFile,
        TokenType::SectionEnd
    };
    automata[TokenType::MatchGroupEnd] = {
		TokenType::ElementDashStart
	};
    automata[TokenType::QuestionEnd] = {
        TokenType::ElementDashStart,
        TokenType::ElementPlusStart
    };
    automata[TokenType::ElementDashStart] = {
        TokenType::IdentifierStart,
        TokenType::TextFragment,
        TokenType::ElementOrderModifier
    };
    automata[TokenType::ElementOrderModifier] = {
		TokenType::TextFragment
    };
    automata[TokenType::ElementPlusStart] = {
        TokenType::TextFragment
    };
    automata[TokenType::Identifier] = {
        TokenType::IdentifierEnd,
        TokenType::SectionStart
    };
    automata[TokenType::IdentifierStart] = {
        TokenType::Identifier
    };
    automata[TokenType::IdentifierEnd] = {
        TokenType::TextFragment
    };
    automata[TokenType::SectionIdentifierStart] = {
        TokenType::Identifier
    };
    automata[TokenType::SectionStart] = {
        TokenType::ElementDashStart,
        TokenType::SectionIdentifierStart,
        TokenType::EndOfFile
    };
    automata[TokenType::SectionEnd] = {
        TokenType::SectionIdentifierStart,
        TokenType::ElementDashStart,
        TokenType::EndOfFile
    };
    automata[TokenType::SOF] = {
        TokenType::ElementDashStart,
        TokenType::SectionIdentifierStart,
        TokenType::EndOfFile
    };
    automata[TokenType::EndOfFile] = {};
    return automata;
}

std::string capitalize(const std::string& str) {
    if (str.empty()) return str;
    std::string result = str;
    result[0] = std::towupper(result[0]);
    return result;
}

Result<NoneType> ValidateGrammar(const std::vector<Token>& tokens) {
    automata = parserAutomata();
    for (size_t i = 0; i < tokens.size() - 1; ++i) {
        Token token = tokens[i];
        Token nextToken = tokens[i + 1];
        if (!contains(automata[token.tokenType], nextToken.tokenType)) {
			return {
				.error=std::format(
					"Invalid token sequence: {} cannot precede {}",
					std::string(capitalize(Token::ToString(&token.tokenType))),
					std::string(capitalize(Token::ToString(&nextToken.tokenType)))
				),
				.row=token.row,
				.column=token.column
			};
        }
    }
	return {};
}

Result<std::vector<Question*>> ParseQuestions(const std::vector<Token>& tokens) {
    auto questions = std::vector<Question*>();

	if (tokens.size() == 0) {
		return {questions, ""};
	}

    auto result = ValidateGrammar(tokens);
    if (result.error.length() > 0) {
		return {
			questions,
			result.error,
			result.row,
			result.column
		};
    }

    std::string section;
    size_t i = 0;

	debug = true;
	if (debug) {
		std::cout << "SECTION: Parser output:\n";
	}
    while (i < tokens.size()) {
        if (tokens[i].tokenType == TokenType::ElementDashStart) {

            std::string id, questionText;
            std::vector<QuestionElement> questionElements;
			bool isOrderQuestion = false;
			bool isGroupQuestion = false;

            // Start element parsing & add to the offset.
            if (tokens[i + 1].tokenType == TokenType::IdentifierStart) {
                id = tokens[i + 2].content;
                questionText = tokens[i + 4].content;
                i += 6;
            } else {
                id = "";
                questionText = tokens[i + 1].content;
                i += 3;
            }

			auto isInBounds = [tokens](size_t i) {
				return i < tokens.size() && tokens[i].tokenType != TokenType::EndOfFile; 
			};

			// Parse elements of a question.
            while (isInBounds(i)) {

				// Handle other constructs.
				if (tokens[i].tokenType == TokenType::SectionIdentifierStart) {
					break;
				}
				if (tokens[i].tokenType == TokenType::SectionEnd) {
					break;
				}

				// Check question end.
                if (tokens[i].tokenType == TokenType::ElementDashStart && isInBounds(i + 3)) {
					// Distance to the possible question end.
					size_t offset = tokens[i + 1].tokenType == TokenType::IdentifierStart ? 5 : 2;
                    if (isInBounds(i + offset) && tokens[i + offset].tokenType == TokenType::QuestionEnd) {
                        break;
                    }
                    if (offset == 5 && tokens[i + 5].tokenType != TokenType::QuestionEnd) {
						// Cannot place the identifier on the ordinary element.
                        return {
							questions,
							"Invalid identifier placement",
							tokens[i].row,
							tokens[i].column
						};
                    }
                }
				
				// Determine element type.
				bool isDash;
				bool isGroup = false;
                if (isInBounds(i+1) && tokens[i].tokenType == TokenType::ElementOrderModifier) {
					isOrderQuestion = true;
					if (!isDash) {
						// TODO: err
					}
				}
                if (tokens[i].tokenType == TokenType::ElementDashStart) {
					isDash = true;
					if (isOrderQuestion) {
						// TODO: err
					}
				} else {
					isDash = false;
				}
				if (isInBounds(i+2) && tokens[i].tokenType == TokenType::MatchGroupEnd) {
					isGroup         = true;
					isGroupQuestion = true;
					if (!isDash) {
						// TODO: err
					}
					if (isOrderQuestion) {
						// TODO: err
					}
				}

                QuestionElement questionElement;
                questionElement.isDash  = isDash;
				questionElement.isGroup = isGroup;
                questionElement.content = tokens[i + 1].content;
                questionElements.push_back(questionElement);

				size_t offset = 2;
				if (isOrderQuestion) {
					offset += 1;
				}
				if (isGroup) {
					offset += 1;
				}

                i += offset;
            }

            if (questionElements.size() > 1) {
				if (isGroupQuestion) {
					GroupQuestion *question = new GroupQuestion();
					// TODO

				} if (isOrderQuestion) {
					auto *question = new MultiElementQuestion();
					question->ID = id;
					question->QuestionText = questionText;
					for (const auto& elem : questionElements) {
						Choice choice;
						choice.Answer = elem.content;
						choice.IsCorrect = !elem.isDash;
						question->Choices.push_back(choice);
					}
					question->Section = section;
					questions.push_back(question);
					if (debug) {
					std::cout << question->ToString() << "\n";
					}
				}
            }
        } else if (tokens[i].tokenType == TokenType::SectionIdentifierStart) {
			if (section.length() > 0) {
				return {
					questions,
					std::format("Section \"{}\" is not closed, cannot open other section", section),
					tokens[i].row,
					tokens[i].column
				};
			}
            section = tokens[i + 1].content;
            i += 3;
			if (debug) {
				std::cout << "Started section: " << section << "\n";
			}
        } else if (tokens[i].tokenType == TokenType::SectionEnd) {
			if (section.length() <= 0) {
				return {
					questions,
					std::format("Cannot close section, no section was opened", section),
					tokens[i].row,
					tokens[i].column
				};
			}
            section.clear();
            i += 1;
			if (debug) {
				std::cout << "Section ended.\n";
			}
        } else if (tokens[i].tokenType == TokenType::EndOfFile) {
			if (debug) {
				std::cout << "File terminated: EndOfFile\n";
			}
            break;
        } else {
            return {
				questions,
				"Unexpected token encountered",
				tokens[i].row,
				tokens[i].column
			};
        }
    }

	if (debug) {
		std::cout << "SECTION END: Parser output:\n";
	}
    return {questions, ""};
}
