Resolver.h

#ifndef DUELSCRIPT_RESOLVER_H
#define DUELSCRIPT_RESOLVER_H

#include "AstNodes.h"
#include <iostream>
#include <vector>
#include <map>
#include <string>
#include <any>

struct VariableInfo {
    std::string name;
    bool defined;
    TokenType type;
    std::string customTypeName;
    int line;
};

struct StructInfo {
    std::string name;
    std::map<std::string, TokenType> fields;
};

class Resolver : public ExprVisitor<std::any>, public StmtVisitor<void> {
public:
    bool hadError = false;

    void resolve(const std::vector<std::unique_ptr<Stmt>>& statements) {
        beginScope();
        for (const auto& stmt : statements) {
            resolve(stmt.get());
        }
        endScope();
    }

private:
    std::vector<std::map<std::string, VariableInfo>> scopes;
    std::map<std::string, StructInfo> structs;
    enum class FunctionType { NONE, RITUAL };
    FunctionType currentFunction = FunctionType::NONE;

    void resolve(Stmt* stmt) { if(stmt) stmt->accept(*this); }

    TokenType resolveExpr(Expr* expr) {
        if(!expr) return TokenType::TOKEN_EOF;
        try {
            return std::any_cast<TokenType>(expr->accept(*this));
        } catch (const std::bad_any_cast&) {
            return TokenType::TOKEN_EOF;
        }
    }

    void beginScope() { scopes.push_back(std::map<std::string, VariableInfo>()); }
    void endScope() { scopes.pop_back(); }

    void declare(Token name, TokenType type, std::string typeName = "") {
        if (scopes.empty()) return;
        std::map<std::string, VariableInfo>& scope = scopes.back();
        if (scope.find(name.lexeme) != scope.end()) {
            error(name, "Variable with this name already declared in this Shadow Realm.");
        }
        scope[name.lexeme] = {name.lexeme, false, type, typeName, name.line};
    }

    void define(Token name) {
        if (scopes.empty()) return;
        scopes.back()[name.lexeme].defined = true;
    }

    TokenType getVariableType(Token name) {
        for (int i = scopes.size() - 1; i >= 0; i--) {
            if (scopes[i].count(name.lexeme)) {
                return scopes[i][name.lexeme].type;
            }
        }
        error(name, "Undeclared variable '" + name.lexeme + "'.");
        return TokenType::TOKEN_EOF;
    }

    // --- (Strict Type Match Check) ---
    void checkTypeMatch(TokenType expected, TokenType actual, int line, const std::string& context) {
        if (expected == TokenType::TOKEN_EOF || actual == TokenType::TOKEN_EOF) return;

        // Strict Check: Types must match EXACTLY. No implicit conversions.
        if (expected == actual) return;

        std::string type1 = tokenTypeToString(expected);
        std::string type2 = tokenTypeToString(actual);

        if (type1.find("KEYWORD_") == 0) type1 = type1.substr(8);
        if (type2.find("KEYWORD_") == 0) type2 = type2.substr(8);

        std::cerr << "[Line " << line << "] Type Error in " << context << ": "
                  << "Type mismatch! Cannot convert " << type2 << " to " << type1 << ". Types must be identical." << std::endl;
        hadError = true;
    }

    void error(Token token, const std::string& message) {
        std::cerr << "[Line " << token.line << "] Semantic Error: " << message << std::endl;
        hadError = true;
    }

    bool isNumber(TokenType type) {
        return type == TokenType::KEYWORD_DARKMAGICIAN || type == TokenType::KEYWORD_BLUEEYESWHITEDRAGON;
    }

public:
    void visitBlockStmt(const BlockStmt& stmt) override {
        beginScope();
        for (const auto& s : stmt.statements) resolve(s.get());
        endScope();
    }

    void visitVarDeclStmt(const VarDeclStmt& stmt) override {
        declare(stmt.name, stmt.type.type, stmt.type.lexeme);
        if (stmt.initializer) {
            TokenType initType = resolveExpr(stmt.initializer.get());
            checkTypeMatch(stmt.type.type, initType, stmt.name.line, "Variable Declaration");
        }
        define(stmt.name);
    }

    void visitFunctionStmt(const FunctionStmt& stmt) override {
        declare(stmt.name, TokenType::KEYWORD_RITUAL);
        define(stmt.name);
        FunctionType enclosing = currentFunction;
        currentFunction = FunctionType::RITUAL;
        beginScope();
        for (const auto& param : stmt.params) {
            declare(param.name, param.type.type, param.type.lexeme);
            define(param.name);
        }
        resolve(stmt.body.get());
        endScope();
        currentFunction = enclosing;
    }

    void visitExpressionStmt(const ExpressionStmt& stmt) override { resolveExpr(stmt.expression.get()); }

    void visitIfStmt(const IfStmt& stmt) override {
        resolveExpr(stmt.condition.get());
        resolve(stmt.thenBranch.get());
        if (stmt.elseBranch) resolve(stmt.elseBranch.get());
    }

    void visitSummonStmt(const SummonStmt& stmt) override { resolveExpr(stmt.expression.get()); }

    void visitDrawStmt(const DrawStmt& stmt) override {
        TokenType varType = getVariableType(stmt.name);
        for (int i = scopes.size() - 1; i >= 0; i--) {
            if (scopes[i].count(stmt.name.lexeme)) {
                scopes[i][stmt.name.lexeme].defined = true;
                return;
            }
        }
    }

    void visitReturnStmt(const ReturnStmt& stmt) override {
        if (currentFunction == FunctionType::NONE) error(stmt.keyword, "Can't 'Tribute' outside a Ritual.");
        if (stmt.value) resolveExpr(stmt.value.get());
    }

    void visitWhileStmt(const WhileStmt& stmt) override {
        resolveExpr(stmt.condition.get());
        resolve(stmt.body.get());
    }

    void visitClassStmt(const ClassStmt& stmt) override {
        declare(stmt.name, TokenType::KEYWORD_LORDOFD);
        define(stmt.name);
        StructInfo info; info.name = stmt.name.lexeme;
        for(const auto& f : stmt.fields) info.fields[f->name.lexeme] = f->type.type;
        structs[stmt.name.lexeme] = info;
    }
    void visitStructStmt(const StructStmt& stmt) override {
        declare(stmt.name, TokenType::KEYWORD_TOONWORLD);
        define(stmt.name);
        StructInfo info; info.name = stmt.name.lexeme;
        for(const auto& f : stmt.fields) info.fields[f->name.lexeme] = f->type.type;
        structs[stmt.name.lexeme] = info;
    }
    void visitIncludeStmt(const IncludeStmt& stmt) override {}
    void visitUsingStmt(const UsingStmt& stmt) override {}

    std::any visitVariableExpr(const VariableExpr& expr) override {
        if (!scopes.empty()) {
            auto& scope = scopes.back();
            auto it = scope.find(expr.name.lexeme);
            if (it != scope.end() && !it->second.defined) {
                error(expr.name, "Can't read local variable in its own initializer.");
            }
        }
        return getVariableType(expr.name);
    }

    std::any visitAssignExpr(const AssignExpr& expr) override {
        TokenType valueType = resolveExpr(expr.value.get());
        TokenType varType = getVariableType(expr.name);
        checkTypeMatch(varType, valueType, expr.name.line, "Assignment");
        return varType;
    }

    std::any visitBinaryExpr(const BinaryExpr& expr) override {
        TokenType left = resolveExpr(expr.left.get());
        TokenType right = resolveExpr(expr.right.get());
        TokenType op = expr.op.type;

        // --- (Strict Binary Operation Checks) ---

        // Arithmetic (+ - * /)
        if (op == TokenType::MINUS || op == TokenType::STAR || op == TokenType::SLASH || op == TokenType::PLUS) {

            // 1. Check if both are numbers (or both strings for +)
            bool isLeftNum = isNumber(left);
            bool isRightNum = isNumber(right);

            if (op == TokenType::PLUS) {
                // Allow string concatenation ONLY if BOTH are strings
                if (left == TokenType::KEYWORD_REDEYESBLACKDRAGON && right == TokenType::KEYWORD_REDEYESBLACKDRAGON) {
                    return TokenType::KEYWORD_REDEYESBLACKDRAGON;
                }
                // Otherwise, must be numbers
                if (!isLeftNum || !isRightNum) {
                    error(expr.op, "Operands for '+' must be strictly same type (two numbers or two strings).");
                    return TokenType::TOKEN_EOF;
                }
            } else {
                // For - * /, must be numbers
                if (!isLeftNum || !isRightNum) {
                    error(expr.op, "Operands must be numbers.");
                    return TokenType::TOKEN_EOF;
                }
            }

            // 2. STRICT TYPE CHECK: Left and Right must be the SAME type
            if (left != right) {
                std::string t1 = tokenTypeToString(left);
                std::string t2 = tokenTypeToString(right);
                error(expr.op, "Type mismatch in binary operation: Cannot operate on different types " + t1 + " and " + t2 + ".");
                return TokenType::TOKEN_EOF;
            }

            // Return the type (since they are same)
            return left;
        }

        // Comparison (> < >= <=)
        if (op == TokenType::GREATER || op == TokenType::GREATER_EQUAL || op == TokenType::LESS || op == TokenType::LESS_EQUAL) {
            if (!isNumber(left) || !isNumber(right)) {
                error(expr.op, "Operands must be numbers for comparison.");
                return TokenType::TOKEN_EOF;
            }
            // Strict check: must be same numeric type
            if (left != right) {
                error(expr.op, "Cannot compare different numeric types.");
                return TokenType::TOKEN_EOF;
            }
            return TokenType::KEYWORD_TIMEWIZARD;
        }

        // Equality (== !=)
        if (op == TokenType::EQUAL_EQUAL || op == TokenType::BANG_EQUAL) {
            if (left != right) {
                error(expr.op, "Cannot check equality between different types.");
            }
            return TokenType::KEYWORD_TIMEWIZARD;
        }

        if (op == TokenType::SUMMON_OP) {
            return TokenType::KEYWORD_REDEYESBLACKDRAGON;
        }

        return TokenType::TOKEN_EOF;
    }

    std::any visitUnaryExpr(const UnaryExpr& expr) override {
        TokenType right = resolveExpr(expr.right.get());
        if (expr.op.type == TokenType::BANG) {
            if (right != TokenType::KEYWORD_TIMEWIZARD) error(expr.op, "Operand for '!' must be TimeWizard.");
            return TokenType::KEYWORD_TIMEWIZARD;
        }
        if (expr.op.type == TokenType::MINUS) {
            if (!isNumber(right)) error(expr.op, "Operand for '-' must be a number.");
            // Result is same type as operand
            return right;
        }
        return TokenType::TOKEN_EOF;
    }

    std::any visitCallExpr(const CallExpr& expr) override {
        resolveExpr(expr.callee.get());
        for(auto& arg : expr.arguments) resolveExpr(arg.get());
        return TokenType::TOKEN_EOF;
    }

    std::any visitGetExpr(const GetExpr& expr) override {
        resolveExpr(expr.object.get());
        return TokenType::TOKEN_EOF;
    }

    std::any visitSetExpr(const SetExpr& expr) override {
        TokenType valType = resolveExpr(expr.value.get());

        if (VariableExpr* var = dynamic_cast<VariableExpr*>(expr.object.get())) {
            for (int i = scopes.size() - 1; i >= 0; i--) {
                if (scopes[i].count(var->name.lexeme)) {
                    VariableInfo& info = scopes[i][var->name.lexeme];
                    if (info.type == TokenType::IDENTIFIER) {
                        if (structs.count(info.customTypeName)) {
                            StructInfo& sInfo = structs[info.customTypeName];
                            if (sInfo.fields.count(expr.name.lexeme)) {
                                TokenType fieldType = sInfo.fields[expr.name.lexeme];
                                // Strict check for struct fields
                                checkTypeMatch(fieldType, valType, expr.name.line, "Struct Field Assignment (" + expr.name.lexeme + ")");
                            }
                        }
                    }
                    break;
                }
            }
        }
        return valType;
    }

    std::any visitGroupingExpr(const GroupingExpr& expr) override { return resolveExpr(expr.expression.get()); }

    std::any visitLiteralExpr(const LiteralExpr& expr) override {
        if (expr.value.type() == typeid(double)) {
            double v = std::any_cast<double>(expr.value);
            if (v == (int)v) return TokenType::KEYWORD_DARKMAGICIAN;
            return TokenType::KEYWORD_BLUEEYESWHITEDRAGON;
        }
        if (expr.value.type() == typeid(std::string)) return TokenType::KEYWORD_REDEYESBLACKDRAGON;
        if (expr.value.type() == typeid(bool)) return TokenType::KEYWORD_TIMEWIZARD;
        return TokenType::TOKEN_EOF;
    }
};

#endif