#include "generator.hpp"
#include <fstream>
#include <iostream>


void Trie::insert(const std::deque<char>& currentKGram) {
    TrieNode* current = root;

    for (char ch : currentKGram) {
        if (current->children.find(ch) == current->children.end()) {
            current->children[ch] = new TrieNode();
        } else { ++current->occurances; }

        current = current->children[ch];
    }

    current->isEndOfWord = true;
}

bool Trie::search(const std::deque<char>& currentKGram) const {
    TrieNode* current = root;

    for (char ch : currentKGram) {
        if (current->children.find(ch) == current->children.end()) {
            return false;
        }

        current = current->children[ch];
    }

    return current->isEndOfWord;
}

void Generator::SetArguments(int argc, char* argv[]) {
    std::string tempStr;
    for (int i = 1; i < argc; i += 2) {
        tempStr.assign(argv[i]);
        if (tempStr == "-i") {
            setup.isFileSet = true;
            setup.filename.assign(argv[i+1]);
        }
        if (tempStr == "-k") {
            setup.isPrefixSet = true;
            setup.prefixLength = std::stoi(argv[i+1]);
        }
        if (tempStr == "-n") {
            setup.isOutputSet = true;
            setup.outputLength = std::stoi(argv[i+1]);
        }
        if (tempStr == "-h") {
            PrintUsage();
            exit(0);
        }
    }
    if (!setup.isFileSet) { std::cerr << "[Setup - Error] Filename not specified" << std::endl; }
    if (!setup.prefixLength) { std::cerr << "[Setup - Error] Prefix length not specified" << std::endl; }
    if (!setup.outputLength) { std::cerr << "[Setup - Error] Output length not specified" << std::endl; }
    if (!setup.isFileSet || !setup.isPrefixSet || !setup.isOutputSet) { PrintUsage(); }
}

void Generator::ReadFile(void) {
    std::ifstream inputFile(setup.filename);
    if (!inputFile.is_open()) {
        std::cerr << "[ReadFile - Error] Could not open file: " << setup.filename << std::endl;
        exit(1);
    }
    std::deque<char> currentKGram;
    char tempChar;
    // Read in first k-gram
    {
        std::string initializeKGram;
        inputFile.get(&initializeKGram[0], setup.prefixLength);
        for (char ch : initializeKGram) { currentKGram.emplace_back(ch); }
        trie.insert(currentKGram);
    }
    // Read rest of file
    while (inputFile.get(tempChar)) {
        currentKGram.emplace_back(tempChar);
        if (currentKGram.size() > setup.prefixLength) { currentKGram.pop_front(); }
        trie.insert(currentKGram);
    }
}

void Generator::GenerateOutput(void) {
}

char GenerateCharacter(void) {
    return 'z';
}

void PrintUsage(void) {
    std::cout << "Usage:  markov -i input_file -k prefix_length -n output_length" << std::endl;
    std::cout << "        -i: Direct path to input file for basis" << std::endl;
    std::cout << "        -k: Prefix length for Markov chain" << std::endl;
    std::cout << "        -n: Length of output to be generated (words)" << std::endl;
    std::cout << "        -h: Prints this usage text" << std::endl;
}