#include <ios>
#include <iostream>
#include <csignal>
#include <fstream>
#include <string>
#include <vector>
#include <variant>
#include <map>
#include <optional>

using namespace std;

enum class valtype {
    INT, DEC, STR, BOOL, KEYWORD, UNKNOWN
};

enum class keywords {
    IF, ELSE, WHILE, INT, DEC, STR, BOOL, FUN, RETURN,
    OPARE, CPARE, OBRAC, CBRAC,
    ADDTO, SUBTRACTFROM, MULTIPLYTO, DIVIDEFROM,
    ADD, SUBTRACT, MULTIPLY, DIVIDE,
    EQUAL, INEQUAL, LESS, GREATER, EQLESS, EQGREATER,
    INCREMENT, DECREMENT,
    PRINT, LET, INPUT, EXIT,
    VALUE, SEMICOLON, VARIABLE
};

struct Value {
    valtype type;
    variant<int, double, string, bool> value;
};

struct Token {
    keywords keyword;
    Value value;
    valtype type = valtype::KEYWORD;
};

struct var {
    valtype type;
    variant<int, double, string, bool> value;

    string toString() const {
        if (type == valtype::INT) {
            return to_string(get<int>(value));
        }
        else if (type == valtype::DEC) {
            return to_string(get<double>(value));
        }
        else if (type == valtype::STR) {
            return get<string>(value);
        }
        else if (type == valtype::BOOL) {
            return to_string(get<bool>(value));
        } else {
            return "unknown";
        }
    }
};

bool debugMode = false;

class ArgParser {
    private:
        vector<string> args;
    public:
        ArgParser(int argc, char* argv[]) {
            // First, collect all arguments
            for (int i = 0; i < argc; i++) {
                args.push_back(argv[i]);
            }
            
            // Then process them
            for (int i = 0; i < args.size(); i++) {
                if (args[i] == "--debug") {
                    debugMode = true;
                    args.erase(args.begin() + i);
                } else if (args[i] == "--help") {
                    cout << "mxlang interpreter" << endl;
                    cout << "Usage: mx [file]" << endl;
                    cout << "Options:" << endl;
                    cout << "    --debug    Enable debug mode" << endl;
                    cout << "    --help     Show this help message" << endl;
                    cout << "Issues? Send an email to max@maxwellj.xyz" << endl;
                    cout << "Report bugs at https://git.maxwellj.xyz/max/mx" << endl;
                    exit(0);
                }
            }
        }
        
        string getArg(int index) {
            if (index >= 0 && index < args.size()) {
                return args[index];
            }
            return "";
        }
};

string logList;

class Logger {
    private:
        bool isDebug = false;

        void writeToLog(string type, string in) {
            logList += type + ": " + in + "\n";
        }

    public:
        string getLog() {
            return logList;
        }
        void toggleDebugPrint() {
            isDebug = !isDebug;
        }
        void error(string in) {
            cout << "Error: " + in + "\n";
            writeToLog("Error", in);
        }
        void fatalError(string in, int code) {
            cout << "Error: " + in + "\n";
            writeToLog("Error", in);
            exit(code);
        }
        void info(string in) {
            cout << "Info: " + in + "\n";
            writeToLog("Info", in);
        }  
        void debug(string in) {
            if (isDebug) cout << "Debug: " + in + "\n";
            writeToLog("Debug", in);
        }
};

class SyntaxError {
    private:
    public:
        void fnTypeMismatch(string function, vector<string> validTypes, valtype typeGiven, string notes = "") {
            cerr << "TypeError: function type mismatch" << endl;
            cerr << "Function '" << function << "' expected one of the following types: ";
            for (int i = 0; i < validTypes.size(); i++) {
                if (i != 0) cerr << ", ";
                cerr << validTypes[i];
            }
            cerr << endl << "Got type '";
            if (typeGiven == valtype::INT) cerr << "int";
            else if (typeGiven == valtype::DEC) cerr << "dec";
            else if (typeGiven == valtype::STR) cerr << "str";
            else if (typeGiven == valtype::BOOL) cerr << "bool";
            else cerr << "unknown";
            cerr << "' instead" << endl;
            if (!notes.empty()) cerr << "Notes: " << notes << endl;
            exit(1);
        }
        void fnNotSufficientArgs(string function, int minArgs, int maxArgs, int argsGiven) {
            cerr << "TypeError: function not sufficient arguments" << endl;
            cerr << "Function '" << function << "' expected between " << minArgs << " and " << maxArgs << " arguments, got " << argsGiven << endl;
            exit(1);
        }
        void unknownFn() {
            cerr << "TypeError: unknown function" << endl;
            exit(1);
        }
        void mathTypeMismatch(string notes = "") {
            cerr << "MathError: cannot add two different types together" << endl;
            if (!notes.empty()) cerr << "Notes: " << notes << endl;
            exit(1);
        }
        void mathTooFewArgs(string notes = "") {
            cerr << "MathError: too few things to add together" << endl;
            if (!notes.empty()) cerr << "Notes: " << notes << endl;
            exit(1);
        }
        void mathCannotDoOperationOnType(string operatorUsed, string type, string notes = "") {
            cerr << "MathError: cannot use operator '" + operatorUsed + " on type '" + type + "'" << endl;
            if (!notes.empty()) cerr << "Notes: " << notes << endl;
            exit(1);
        }
        void generalError(string notes = "") {
            cerr << "GeneralError: Something went awfully wrong and we don't know why lmao" << endl;
            if (!notes.empty()) cerr << "Notes: " << notes << endl;
            exit(1);
        }
};

class Parser {
    private:
        string buffer;
        vector<string> lines;
        vector<string> termBuffer;
        vector<Token> tokens;
        vector<vector<string>> terms;
        vector<vector<keywords>> things;
        vector<vector<variant<var, keywords>>> stuff;
        bool canInt(string in) {
            try {
                stoi(in);
                return true;
            } catch (...) {
                return false;
            }
        }
        bool canDec(string in) {
            try {
                stod(in);
                return true;
            } catch (...) {
                return false;
            }
        }
    public:
        void parseLines(string in, Logger log) {
            log.debug("Parsing lines...");
            buffer.clear();
            termBuffer.clear();
            terms.clear();
            bool isString = false;
            bool isEscaped = false;

            for (size_t i = 0; i < in.size(); i++) {
                char c = in[i];
                
                if (isEscaped) {
                    buffer += c;
                    isEscaped = false;
                    continue;
                }

                if (c == '\\') {
                    isEscaped = true;
                    buffer += c;
                    continue;
                }

                if (c == '"') {
                    buffer += c;
                    isString = !isString;
                } else if (c == '\n' && !isString) {
                    // Skip newlines outside strings
                    continue;
                } else if (c == ';' && !isString) {
                    if (!buffer.empty()) {
                        termBuffer.push_back(buffer);
                        buffer.clear();
                    }
                    terms.push_back(termBuffer);
                    termBuffer.clear();
                } else if ((c == ',' || c == '{' || c == '}' || c == '(' || c == ')') && !isString) {
                    if (!buffer.empty()) {
                        termBuffer.push_back(buffer);
                        buffer.clear();
                    }
                    termBuffer.push_back(string(1, c));
                } else if (c == ' ' && !isString) {
                    if (!buffer.empty()) {
                        termBuffer.push_back(buffer);
                        buffer.clear();
                    }
                } else {
                    buffer += c;
                }
            }

            // Handle any remaining buffer content
            if (!buffer.empty()) {
                termBuffer.push_back(buffer);
            }
            if (!termBuffer.empty()) {
                terms.push_back(termBuffer);
            }

            // Debug output
            string debugString;
            log.debug("Lines have been parsed.");
            for (const auto& term : terms) {
                for (const auto& t : term) {
                    debugString += t + ", ";
                }
                debugString += "\n";
            }
            log.debug(debugString);
        }
        void processLines() {
            // Process vector<vector<string>> terms
            for (int i = 0; i < terms.size(); i++) {
                for (int j = 0; j < terms[i].size(); j++) {
                    Token token;
                    string ct = terms[i][j];
                    if (ct == " ") continue;
                    // Oh boy here we go
                    else if (ct == "fun") token.keyword = keywords::FUN;
                    else if (ct == "let") token.keyword = keywords::LET;
                    else if (ct == "print") token.keyword = keywords::PRINT;
                    else if (ct == "return") token.keyword = keywords::RETURN;
                    else if (ct == "exit") token.keyword = keywords::EXIT;
                    else if (ct == "int") token.keyword = keywords::INT;
                    else if (ct == "dec") token.keyword = keywords::DEC;
                    else if (ct == "str") token.keyword = keywords::STR;
                    else if (ct == "bool") token.keyword = keywords::BOOL;
                    else if (ct == "{") token.keyword = keywords::OBRAC;
                    else if (ct == "}") token.keyword = keywords::CBRAC;
                    else if (ct == "(") token.keyword = keywords::OPARE;
                    else if (ct == ")") token.keyword = keywords::CPARE;
                    else if (ct == "+") token.keyword = keywords::ADD;
                    else if (ct == "-") token.keyword = keywords::SUBTRACT;
                    else if (ct == "*") token.keyword = keywords::MULTIPLY;
                    else if (ct == "/") token.keyword = keywords::DIVIDE;
                    else if (ct == "++") token.keyword = keywords::INCREMENT;
                    else if (ct == "--") token.keyword = keywords::DECREMENT;
                    else if (ct == "+=") token.keyword = keywords::ADDTO;
                    else if (ct == "-=") token.keyword = keywords::SUBTRACTFROM;
                    else if (ct == "*=") token.keyword = keywords::MULTIPLYTO;
                    else if (ct == "/=") token.keyword = keywords::DIVIDEFROM;
                    else {
                        token.keyword = keywords::VALUE;
                        // Convert the value based on its type
                        if (canDec(ct)) {
                            token.type = valtype::DEC;
                            token.value.type = valtype::DEC;
                            token.value.value = stod(ct);
                        }
                        else if (canInt(ct)) {
                            token.type = valtype::INT;
                            token.value.type = valtype::INT;
                            token.value.value = stoi(ct);
                        }
                        else if (ct == "true" || ct == "false") {
                            token.type = valtype::BOOL;
                            token.value.type = valtype::BOOL;
                            token.value.value = (ct == "true");
                        }
                        else {
                            // Handle strings - remove quotes if present
                            token.type = valtype::STR;
                            token.value.type = valtype::STR;
                            if (ct.size() >= 2 && ct.front() == '"' && ct.back() == '"') {
                                // Remove the quotes
                                token.value.value = ct.substr(1, ct.size() - 2);
                            } else {
                                token.value.value = ct;
                            }
                        }
                    }
                    tokens.push_back(token);
                }
                Token semi;
                semi.keyword = keywords::SEMICOLON;
                tokens.push_back(semi);
            }
}
        vector<Token> getTokens() {
            return tokens;
        }
};

class Interpreter {
    private:
        SyntaxError syntaxError;
        vector<Token> tokens;
        map<string, Value> variables;
        Logger log;
        int tokenIndex = -1;
        optional<Token> consume() {
            tokenIndex++;
            if (tokenIndex < tokens.size()) return tokens[tokenIndex];
            return {};
        } 
        optional<Token> peek(int offset = 1) {
            int index = tokenIndex + offset;
            if (index >= 0 && index < tokens.size()) return tokens[index];
            return {};
        }

    public:
        void convertToTokens(vector<Token> tokenList) {
            if (debugMode) log.toggleDebugPrint();
            tokens = tokenList;
            log.debug("Alright we got " + to_string(tokens.size()) + " tokens");
            
            while (tokenIndex < static_cast<int>(tokens.size() - 1)) {
                auto currentToken = consume();
                if (!currentToken) break;
                
                vector<Token> currentInstruction;
                currentInstruction.push_back(currentToken.value());
                
                // Collect tokens until semicolon
                while (auto nextToken = peek(1)) {
                    if (nextToken->keyword == keywords::SEMICOLON) {
                        consume(); // consume the semicolon
                        break;
                    }
                    consume(); // consume the peeked token
                    currentInstruction.push_back(nextToken.value());
                }
                // Apply variables to tokens
                // We start at 1 so we can reassign variables in the execution of code
                for (int i = 1; i < currentInstruction.size(); i++) {
                    if (currentInstruction[i].type == valtype::STR) {
                        string potentialVarName = get<string>(currentInstruction[i].value.value);
                        auto varIt = variables.find(potentialVarName);
                        
                        if (varIt != variables.end()) {
                            // Replace the token with the variable's value
                            Token newToken;
                            newToken.keyword = keywords::VALUE;
                            newToken.type = varIt->second.type;
                            newToken.value = varIt->second;
                            currentInstruction[i] = newToken;
                        }
                    }
                }
                // Do math
                for (int i = 0; i < currentInstruction.size(); i++) {
                    if (currentInstruction[i].keyword == keywords::ADD || currentInstruction[i].keyword == keywords::SUBTRACT || currentInstruction[i].keyword == keywords::MULTIPLY || currentInstruction[i].keyword == keywords::DIVIDE) {
                        Token newToken;
                        newToken.keyword = keywords::VALUE;
                        if (currentInstruction.size() < i + 1) syntaxError.mathTooFewArgs();
                        Token before = currentInstruction[i - 1];
                        Token after = currentInstruction[i + 1];
                        if (before.type != after.type) syntaxError.mathTypeMismatch();
                        newToken.type = before.type;
                        if (currentInstruction[i].keyword == keywords::ADD) {
                            if (newToken.type == valtype::INT) {
                                newToken.value.value = get<int>(before.value.value) + get<int>(after.value.value);
                            }
                            else if (newToken.type == valtype::DEC) {
                                newToken.value.value = get<double>(before.value.value) + get<double>(after.value.value);
                            }
                            else if (newToken.type == valtype::STR) {
                                newToken.value.value = get<string>(before.value.value) + get<string>(after.value.value);
                            } else {
                                syntaxError.mathCannotDoOperationOnType("+", "bool");
                            }
                        }
                        else if (currentInstruction[i].keyword == keywords::SUBTRACT) {
                            if (newToken.type == valtype::INT) {
                                newToken.value.value = get<int>(before.value.value) - get<int>(after.value.value);
                            }
                            else if (newToken.type == valtype::DEC) {
                                newToken.value.value = get<double>(before.value.value) - get<double>(after.value.value);
                            } else {
                                syntaxError.mathCannotDoOperationOnType("-", "bool or string");
                            }
                        }
                        else if (currentInstruction[i].keyword == keywords::MULTIPLY) {
                            if (newToken.type == valtype::INT) {
                                newToken.value.value = get<int>(before.value.value) * get<int>(after.value.value);
                            }
                            else if (newToken.type == valtype::DEC) {
                                newToken.value.value = get<double>(before.value.value) * get<double>(after.value.value);
                            } else {
                                syntaxError.mathCannotDoOperationOnType("*", "bool or string");
                            }
                        }
                        else if (currentInstruction[i].keyword == keywords::DIVIDE) {
                            if (newToken.type == valtype::INT) {
                                newToken.value.value = get<int>(before.value.value) / get<int>(after.value.value);
                            }
                            else if (newToken.type == valtype::DEC) {
                                newToken.value.value = get<double>(before.value.value) / get<double>(after.value.value);
                            } else {
                                syntaxError.mathCannotDoOperationOnType("/", "bool or string");
                            }
                        } else {
                            // Something has gone terribly wrong
                            // We should never reach this point in the code
                            syntaxError.generalError("The math aint mathing");
                        }
                        // Insert our cool new token and get rid of the boring old stuff
                        currentInstruction[i - 1] = newToken;
                        currentInstruction.erase(currentInstruction.begin() + i);
                        currentInstruction.erase(currentInstruction.begin() + i);
                        i -= 1;
                    }
                }
                // Execute the instruction
                executeCode(currentInstruction);
            }
        }
        void executeCode(vector<Token> tokens) {
            SyntaxError syntaxError;
            log.debug("Token length for this expression is " + to_string(tokens.size()));
            for (int i = 0; i < tokens.size(); i++) {
                if (tokens[i].keyword == keywords::PRINT) {
                    i++;
                    if (tokens.size() <= i) break;
                    Token nextToken = tokens[i];
                    // Handle different value types
                    if (nextToken.keyword == keywords::VALUE) {
                        switch (nextToken.type) {
                            case valtype::INT:
                                cout << get<int>(nextToken.value.value) << endl;
                                break;
                            case valtype::DEC:
                                cout << get<double>(nextToken.value.value) << endl;
                                break;
                            case valtype::STR:
                                cout << get<string>(nextToken.value.value) << endl;
                                break;
                            case valtype::BOOL:
                                cout << (get<bool>(nextToken.value.value) ? "true" : "false") << endl;
                                break;
                            default:
                                vector<string> validTypes = {"int", "dec", "str", "bool"};
                                syntaxError.fnTypeMismatch("print", validTypes, nextToken.type);
                        }
                    } else {
                        syntaxError.fnNotSufficientArgs("print", 1, 1, 0);
                    }
                } else if (tokens[i].keyword == keywords::EXIT) {
                    i++;
                    if (tokens.size() <= i) break;
                    Token nextToken = tokens[i];
                    if (nextToken.keyword == keywords::VALUE) {
                        switch (nextToken.type) {
                            case valtype::INT:
                                exit(get<int>(nextToken.value.value));
                                break;
                            case valtype::DEC:
                                exit(get<double>(nextToken.value.value));
                                break;
                            default:
                                vector<string> validTypes = {"int", "dec"};
                                syntaxError.fnTypeMismatch("exit", validTypes, nextToken.type);
                        }
                    }
                } else if (tokens[i].keyword == keywords::LET) {
                    i++;
                    if (tokens.size() <= i + 2) {
                        syntaxError.fnNotSufficientArgs("let", 3, 3, tokens.size() - i);
                        break;
                    }
                    
                    Token typeToken = tokens[i];
                    Token nameToken = tokens[i + 1];
                    Token valueToken = tokens[i + 2];

                    i += 2;
                    
                    // Validate that we have a valid variable name
                    if (nameToken.type != valtype::STR) {
                        vector<string> validTypes = {"str"};
                        syntaxError.fnTypeMismatch("let (variable name)", validTypes, nameToken.type);
                        continue;
                    }
                    
                    string varName = get<string>(nameToken.value.value);
                    Value newValue;
                    
                    // Check the type declaration matches the value
                    if (typeToken.keyword == keywords::INT && valueToken.type == valtype::INT) {
                        newValue.type = valtype::INT;
                        newValue.value = get<int>(valueToken.value.value);
                    }
                    else if (typeToken.keyword == keywords::DEC && valueToken.type == valtype::DEC) {
                        newValue.type = valtype::DEC;
                        newValue.value = get<double>(valueToken.value.value);
                    }
                    else if (typeToken.keyword == keywords::STR && valueToken.type == valtype::STR) {
                        newValue.type = valtype::STR;
                        newValue.value = get<string>(valueToken.value.value);
                    }
                    else if (typeToken.keyword == keywords::BOOL && valueToken.type == valtype::BOOL) {
                        newValue.type = valtype::BOOL;
                        newValue.value = get<bool>(valueToken.value.value);
                    }
                    else {
                        vector<string> validTypes;
                        if (typeToken.keyword == keywords::INT) validTypes = {"int"};
                        else if (typeToken.keyword == keywords::DEC) validTypes = {"dec"};
                        else if (typeToken.keyword == keywords::STR) validTypes = {"str"};
                        else if (typeToken.keyword == keywords::BOOL) validTypes = {"bool"};
                        syntaxError.fnTypeMismatch("let", validTypes, valueToken.type, "Variable name is " + varName);
                        continue;
                    }
                    
                    // Store the variable
                    variables[varName] = newValue;
                } else {
                    if (tokens[i].keyword == keywords::VALUE && tokens[i].value.type == valtype::STR && variables.find(get<string>(tokens[i].value.value)) != variables.end()) {
                        log.debug("Manipulating variable...");
                        if (tokens.size() < i + 1) {
                            syntaxError.mathTooFewArgs();
                        }
                        string varName = get<string>(tokens[i].value.value);
                        i++;
                        if (variables[varName].type == valtype::INT) {
                            if (tokens[i].keyword == keywords::INCREMENT) {
                                variables[varName].value = get<int>(variables[varName].value) + 1;
                            } else if (tokens[i].keyword == keywords::DECREMENT) {
                                variables[varName].value = get<int>(variables[varName].value) - 1;
                            } else if (tokens[i].keyword == keywords::ADDTO) {
                                if (tokens.size() < i + 1) syntaxError.mathTooFewArgs();
                                if (tokens[i + 1].value.type != valtype::INT) syntaxError.mathTypeMismatch("Expected an int when adding to an int");
                                variables[varName].value = get<int>(variables[varName].value) + get<int>(tokens[i + 1].value.value);
                            } else if (tokens[i].keyword == keywords::SUBTRACTFROM) {
                                if (tokens.size() < i + 1) syntaxError.mathTooFewArgs();
                                if (tokens[i + 1].value.type != valtype::INT) syntaxError.mathTypeMismatch("Expected an int when subtracting from an int");
                                variables[varName].value = get<int>(variables[varName].value) - get<int>(tokens[i + 1].value.value);
                            }
                        } else if (variables[varName].type == valtype::DEC) {
                            if (tokens[i].keyword == keywords::INCREMENT) {
                                variables[varName].value = get<double>(variables[varName].value) + 1;
                            } else if (tokens[i].keyword == keywords::DECREMENT) {
                                variables[varName].value = get<double>(variables[varName].value) - 1;
                            } else if (tokens[i].keyword == keywords::ADDTO) {
                                if (tokens.size() < i + 1) syntaxError.mathTooFewArgs();
                                if (tokens[i + 1].value.type != valtype::DEC) syntaxError.mathTypeMismatch("Expected an dec when adding to an dec");
                                variables[varName].value = get<double>(variables[varName].value) + get<double>(tokens[i + 1].value.value);
                            } else if (tokens[i].keyword == keywords::SUBTRACTFROM) {
                                if (tokens.size() < i + 1) syntaxError.mathTooFewArgs();
                                if (tokens[i + 1].value.type != valtype::DEC) syntaxError.mathTypeMismatch("Expected an dec when subtracting from an dec");
                                variables[varName].value = get<double>(variables[varName].value) - get<double>(tokens[i + 1].value.value);
                            }
                        }
                    } else {
                        syntaxError.unknownFn();
                    }
                }
            }
        }
};

int main(int argc, char* argv[]) {
    // Parse and act upon command line arguments
    ArgParser args(argc, argv);
    // Initialise the logger
    Logger log;
    if (debugMode) log.toggleDebugPrint();
    log.debug("Logger initialised!");
    // Initialise the parser
    Parser parser;
    // Exit if file doesn't exist
    if (argc < 2) log.fatalError("Please provide a file", 1);
    // Read the file
    ifstream inputFile(args.getArg(1));
    string input;
    string file;
    while(getline(inputFile, input)) {
        file += input;
    }
    inputFile.close();
    // Parse the file
    parser.parseLines(file, log);
    parser.processLines();
    // Initialise the interpreter
    Interpreter interpreter;
    // Convert to tokens and run the code
    interpreter.convertToTokens(parser.getTokens());
    return 0;
}