main.cpp

#include <iostream>
#include <string>
#include <vector>
#include <unordered_set>
#include <sstream>

#include "lexer.h"
#include "parser.h"
#include "interpreter.h"

// ===== ANSI COLORS =====
static const char* C_RESET  = "\033[0m";
static const char* C_BOLD   = "\033[1m";
static const char* C_GRAY   = "\033[90m";
static const char* C_GREEN  = "\033[32m";
static const char* C_CYAN   = "\033[36m";
static const char* C_RED    = "\033[31m";

// highlight colors
static const char* K_KEYWORD = "\033[34m"; // blue
static const char* K_BOOL    = "\033[35m"; // purple

// ===== keyword sets (ONLY these get colored) =====
static const std::unordered_set<std::string> KEYWORDS = {
    "int","float","bool","string","char",
    "fun","done","return",
    "while","for",
    "if","else","do",
    "print","input"
};
static const std::unordered_set<std::string> BOOLS = {"true","false"};

static std::string trim(const std::string& s) {
    size_t a = 0;
    while (a < s.size() && isspace((unsigned char)s[a])) a++;
    size_t b = s.size();
    while (b > a && isspace((unsigned char)s[b-1])) b--;
    return s.substr(a, b - a);
}
static bool isAllWhitespace(const std::string& s) {
    return trim(s).empty();
}

// ===== highlight ONLY keywords/bools, do NOT color strings/chars/numbers =====
static std::string highlightKeywordsOnly(const std::string& line) {
    std::ostringstream out;
    std::string tok;

    bool inString = false;
    bool inChar = false;

    auto flushTok = [&]() {
        if (tok.empty()) return;
        if (KEYWORDS.count(tok)) out << K_KEYWORD << tok << C_RESET;
        else if (BOOLS.count(tok)) out << K_BOOL << tok << C_RESET;
        else out << tok;
        tok.clear();
    };

    for (size_t i = 0; i < line.size(); ++i) {
        char c = line[i];

        // keep strings/chars untouched
        if (c == '"' && !inChar) {
            flushTok();
            out << c;
            inString = !inString;
            continue;
        }
        if (c == '\'' && !inString) {
            flushTok();
            out << c;
            inChar = !inChar;
            continue;
        }
        if (inString || inChar) {
            out << c;
            continue;
        }

        if (isalnum((unsigned char)c) || c == '_') {
            tok.push_back(c);
        } else {
            flushTok();
            out << c;
        }
    }
    flushTok();
    return out.str();
}

// ===== completion detection for multi-line =====
struct CompletionState {
    bool ok = true;
    std::string errorMsg;

    int paren = 0;
    int blockDepth = 0;  // blocks closed by done;
    int funcDepth  = 0;  // fun blocks closed by done - return ...;

    bool pendingBlock = false;
    bool pendingFunc  = false;

    bool complete = false;
};

static CompletionState analyzeCompletion(const std::vector<Token>& tks) {
    CompletionState st;
    TokenType last = TokenType::INVALID;

    for (size_t i = 0; i < tks.size(); ++i) {
        const Token& tk = tks[i];

        if (tk.type == TokenType::INVALID) {
            st.ok = false;
            st.errorMsg = tk.lexeme;
            return st;
        }

        if (tk.type != TokenType::END) last = tk.type;

        switch (tk.type) {
            case TokenType::LPAREN: st.paren++; break;
            case TokenType::RPAREN: st.paren--; break;

            // start of a block: while, for, if
            case TokenType::KW_WHILE:
            case TokenType::KW_FOR:
            case TokenType::KW_IF:
                st.pendingBlock = true;
                break;

            // else closes the previous then-block and opens a new block
            case TokenType::KW_ELSE: {
                if (st.blockDepth > 0) st.blockDepth--;
                st.pendingBlock = true;
                break;
            }

            // function
            case TokenType::KW_FUN:
                st.funcDepth++;
                st.pendingFunc = true;
                break;

            case TokenType::COLON:
                if (st.pendingBlock) { st.blockDepth++; st.pendingBlock = false; }
                if (st.pendingFunc)  { st.pendingFunc  = false; }
                break;

            case TokenType::KW_DONE: {
                TokenType next = (i + 1 < tks.size()) ? tks[i + 1].type : TokenType::END;
                if (next == TokenType::MINUS) {
                    if (st.funcDepth > 0) st.funcDepth--;
                } else {
                    if (st.blockDepth > 0) st.blockDepth--;
                }
                break;
            }

            default:
                break;
        }

        if (st.paren < 0) {
            st.ok = false;
            st.errorMsg = "Too many ')'";
            return st;
        }
    }

    bool endsWithSemi = (last == TokenType::SEMI);

    st.complete =
        st.ok &&
        st.paren == 0 &&
        st.blockDepth == 0 &&
        st.funcDepth == 0 &&
        !st.pendingBlock &&
        !st.pendingFunc &&
        endsWithSemi;

    return st;
}

static void clearScreen() {
    std::cout << "\033[2J\033[H" << std::flush;
}

static void printHelp() {
    std::cout
        << C_BOLD << "Qarebaq REPL commands" << C_RESET << "\n"
        << "  " << C_CYAN << ":help"  << C_RESET << "   show help\n"
        << "  " << C_CYAN << ":reset" << C_RESET << "  clear current multi-line buffer\n"
        << "  " << C_CYAN << ":clear" << C_RESET << "  clear screen\n"
        << "  " << C_CYAN << ":exit"  << C_RESET << "   quit\n\n";
}

// rewrite the last typed line with colors (no duplicate printing)
static void rewriteLastInputLine(const std::string& prompt, const std::string& rawLine, bool enableColor) {
    // Move cursor up one line, clear it, and print colored version
    std::cout << "\033[A\r\033[2K";
    if (!enableColor) {
        std::cout << prompt << rawLine << "\n";
        return;
    }
    std::cout << prompt << highlightKeywordsOnly(rawLine) << "\n";
}

int main() {
    std::cout << C_BOLD << "Qarebaq REPL" << C_RESET
              << "  (" << C_GRAY << "type :help" << C_RESET << ")\n";

    Interpreter interp;

    // keep programs alive (important for storing function pointers)
    std::vector<std::vector<StmtPtr>> allPrograms;

    std::string buffer;

    // allow turning colors off if needed
    bool colorOn = true;

    while (true) {
        bool multiline = !buffer.empty();

        std::string prompt = multiline ? std::string(C_GREEN) +C_BOLD+ "....> " + C_RESET
                                       : std::string(C_RED)   +C_BOLD+ ">>> "   + C_RESET;

        // print prompt
        std::cout << prompt << std::flush;

        std::string line;
        if (!std::getline(std::cin, line)) break;

        std::string t = trim(line);

        // Commands
        if (!multiline) {
            if (t == ":exit" || t == "exit") break;
            if (t == ":help") { printHelp(); continue; }
            if (t == ":clear") { clearScreen(); continue; }
            if (t == ":reset") { buffer.clear(); continue; }
            if (t == ":color on") { colorOn = true; continue; }
            if (t == ":color off") { colorOn = false; continue; }
        } else {
            if (t == ":reset") { buffer.clear(); continue; }
            if (t == ":clear") { clearScreen(); continue; }
            if (t == ":color on") { colorOn = true; continue; }
            if (t == ":color off") { colorOn = false; continue; }
        }

        // ignore empty lines only when not in a block
        if (!multiline && isAllWhitespace(line)) {
            // optionally remove the empty visual line by rewriting it to just prompt
            continue;
        }

        // rewrite the line user just typed (no duplicate echo)
        rewriteLastInputLine(prompt, line, colorOn);

        buffer += line + "\n";

        Lexer lx(buffer);
        auto tokens = lx.tokenize();
        CompletionState st = analyzeCompletion(tokens);

        if (!st.ok) {
            std::cout << C_RED << "Lexer error: " << st.errorMsg << C_RESET << "\n";
            buffer.clear();
            continue;
        }

        if (!st.complete) {
            continue; // keep collecting lines
        }

        try {
            Parser ps(tokens);
            auto program = ps.parseProgram();
            allPrograms.push_back(std::move(program));
            interp.run(allPrograms.back());
        } catch (const std::exception& e) {
            std::cout << C_RED << "Error: " << e.what() << C_RESET << "\n";
        }

        buffer.clear();
    }

    std::cout << C_GRAY << "bye.\n" << C_RESET;
    return 0;
}