#include "./builder.hpp"

namespace lsp::language::symbol
{
    Builder::Builder(SymbolTable& table) :
        table_(table), current_scope_id_(kInvalidScopeId)
    {
    }

    void Builder::Build(ast::ASTNode& root)
    {
        current_scope_id_ = table_.CreateScope(
            ScopeKind::Global,
            root.span,
            std::nullopt,
            std::nullopt);

        SymbolId global_symbol = table_.CreateSymbol(
            "::",
            SymbolKind::Namespace,
            root.span,
            current_scope_id_,
            std::nullopt,
            std::nullopt);

        table_.GetScopeManager().AddSymbol(current_scope_id_, "::", global_symbol);

        root.Accept(*this);
    }

    SymbolId Builder::CreateSymbol(const std::string& name, SymbolKind kind, const ast::ASTNode& node, const std::optional<std::string>& type_hint, bool is_class_method)
    {
        return table_.CreateSymbol(
            name,
            kind,
            node.span,
            current_scope_id_,
            current_parent_symbol_id_,
            type_hint,
            is_class_method);
    }

    ScopeId Builder::EnterScopeWithSymbol(ScopeKind kind, SymbolId symbol_id, const ast::Location& range)
    {
        ScopeId parent_scope = current_scope_id_;
        current_scope_id_ = table_.CreateScope(kind, range, parent_scope, symbol_id);
        return current_scope_id_;
    }

    ScopeId Builder::EnterScope(ScopeKind kind, const ast::Location& range)
    {
        ScopeId parent_scope = current_scope_id_;
        current_scope_id_ = table_.CreateScope(kind, range, parent_scope, std::nullopt);
        return current_scope_id_;
    }

    void Builder::ExitScope()
    {
        auto info = table_.GetScopeManager().GetScopeInfo(current_scope_id_);
        if (info && info->parent_scope_id)
            current_scope_id_ = *info->parent_scope_id;
    }

    void Builder::VisitStatements(const std::vector<ast::StatementPtr>& statements)
    {
        for (const auto& stmt : statements)
        {
            if (stmt)
                stmt->Accept(*this);
        }
    }

    void Builder::VisitExpression(ast::Expression& expr)
    {
        expr.Accept(*this);
    }

    std::string Builder::FormatSignature(const ast::Signature& signature)
    {
        std::string result = "(";
        for (size_t i = 0; i < signature.parameters.size(); ++i)
        {
            if (i > 0)
                result += ", ";
            const auto& param = signature.parameters[i];

            if (param.is_var)
                result += "var ";
            if (param.is_out)
                result += "out ";

            result += param.name;
            if (param.type)
                result += ": " + param.type->name;
        }
        result += ")";

        if (signature.return_type)
            result += ": " + signature.return_type->name;

        return result;
    }

    void Builder::ProcessLeftHandSide(const ast::LeftHandSide& lhs, bool is_write)
    {
        std::visit([this, is_write](auto& variant) {
            using T = std::decay_t<decltype(variant)>;

            if constexpr (std::is_same_v<T, std::unique_ptr<ast::Identifier>>)
            {
                if (variant)
                {
                    auto symbol_id = table_.GetScopeManager().FindInScopeChain(current_scope_id_, variant->name);
                    if (symbol_id)
                        table_.AddReference(*symbol_id, variant->span, is_write);
                }
            }
            else if constexpr (std::is_same_v<T, std::unique_ptr<ast::AttributeExpression>>)
            {
                if (variant)
                    VisitExpression(*variant->object);
            }
            else if constexpr (std::is_same_v<T, std::unique_ptr<ast::SubscriptExpression>>)
            {
                if (variant)
                    VisitExpression(*variant->value);
            }
        },
                   lhs);
    }

    void Builder::VisitProgram(ast::Program& node)
    {
        VisitStatements(node.statements);
    }

    void Builder::VisitUnitDefinition(ast::UnitDefinition& node)
    {
        auto symbol_id = CreateSymbol(node.name, SymbolKind::Module, node);

        EnterScopeWithSymbol(ScopeKind::Unit, symbol_id, node.span);
        auto prev_parent = current_parent_symbol_id_;
        current_parent_symbol_id_ = symbol_id;

        VisitStatements(node.interface_statements);
        VisitStatements(node.implementation_statements);
        VisitStatements(node.initialization_statements);
        VisitStatements(node.finalization_statements);

        current_parent_symbol_id_ = prev_parent;
        ExitScope();
    }

    void Builder::VisitClassDefinition(ast::ClassDefinition& node)
    {
        auto symbol_id = CreateSymbol(node.name, SymbolKind::Class, node);

        // 处理继承关系
        for (const auto& parent_class : node.parent_classes)
        {
            auto parent_symbols = table_.GetDefinitionStore().FindByName(parent_class.name);
            for (const auto* parent_def : parent_symbols)
            {
                if (parent_def->kind == SymbolKind::Class)
                    table_.GetInheritanceGraph().AddInheritance(symbol_id, parent_def->id);
            }
        }

        EnterScopeWithSymbol(ScopeKind::Class, symbol_id, node.span);
        auto prev_parent = current_parent_symbol_id_;
        current_parent_symbol_id_ = symbol_id;

        for (const auto& member : node.members)
        {
            if (member)
                member->Accept(*this);
        }

        current_parent_symbol_id_ = prev_parent;
        ExitScope();
    }

    void Builder::VisitFunctionDefinition(ast::FunctionDefinition& node)
    {
        auto symbol_id = CreateSymbol(node.name, SymbolKind::Function, node);

        // 使用新的 Update 接口更新符号信息
        table_.GetDefinitionStore().Update(symbol_id, [&](SymbolDefinition& def) {
            def.detail = FormatSignature(node.signature);
        });

        EnterScopeWithSymbol(ScopeKind::Function, symbol_id, node.span);
        auto prev_function = current_function_id_;
        auto prev_parent = current_parent_symbol_id_;

        current_function_id_ = symbol_id;
        current_parent_symbol_id_ = symbol_id;

        // 添加参数
        for (const auto& param : node.signature.parameters)
        {
            std::optional<std::string> type_hint;
            if (param.type)
                type_hint = param.type->name;
            CreateSymbol(param.name, SymbolKind::Variable, node, type_hint);
        }

        current_parent_symbol_id_ = prev_parent;

        if (node.body)
            node.body->Accept(*this);

        current_function_id_ = prev_function;
        ExitScope();
    }

    void Builder::VisitFunctionDeclaration(ast::FunctionDeclaration& node)
    {
        auto symbol_id = CreateSymbol(node.name, SymbolKind::Function, node);

        table_.GetDefinitionStore().Update(symbol_id, [&](SymbolDefinition& def) {
            def.detail = FormatSignature(node.signature);
        });
    }

    void Builder::VisitMethodDeclaration(ast::MethodDeclaration& node)
    {
        SymbolKind kind = node.method_type == ast::MethodType::kConstructor ? SymbolKind::Constructor : SymbolKind::Method;

        auto symbol_id = CreateSymbol(node.name, kind, node, "", node.is_class_method);

        table_.GetDefinitionStore().Update(symbol_id, [&](SymbolDefinition& def) {
            def.detail = FormatSignature(node.signature);
            def.method_modifier = node.modifier;
        });

        EnterScopeWithSymbol(ScopeKind::Function, symbol_id, node.span);
        auto prev_function = current_function_id_;
        auto prev_parent = current_parent_symbol_id_;

        current_function_id_ = symbol_id;

        for (const auto& param : node.signature.parameters)
        {
            std::optional<std::string> type_hint;
            if (param.type)
                type_hint = param.type->name;
            auto param_prev = current_parent_symbol_id_;
            current_parent_symbol_id_ = symbol_id;
            CreateSymbol(param.name, SymbolKind::Variable, node, type_hint);
            current_parent_symbol_id_ = param_prev;
        }
        if (node.body)
            node.body->Accept(*this);

        current_parent_symbol_id_ = prev_parent;
        current_function_id_ = prev_function;
        ExitScope();
    }

    void Builder::VisitPropertyDeclaration(ast::PropertyDeclaration& node)
    {
        CreateSymbol(node.name, SymbolKind::Property, node, node.type_name);
    }

    void Builder::VisitExternalMethodDefinition(ast::ExternalMethodDefinition& node)
    {
        SymbolKind kind = node.method_type == ast::MethodType::kConstructor ? SymbolKind::Constructor : SymbolKind::Method;

        auto symbol_id = CreateSymbol(node.name, kind, node, std::nullopt, node.is_class_method);

        table_.GetDefinitionStore().Update(symbol_id, [&](SymbolDefinition& def) {
            def.detail = FormatSignature(node.signature);
            def.method_modifier = node.modifier;
        });

        EnterScopeWithSymbol(ScopeKind::Function, symbol_id, node.span);
        auto prev_function = current_function_id_;
        auto prev_parent = current_parent_symbol_id_;

        current_function_id_ = symbol_id;

        for (const auto& param : node.signature.parameters)
        {
            std::optional<std::string> type_hint;
            if (param.type)
                type_hint = param.type->name;
            auto param_prev = current_parent_symbol_id_;
            current_parent_symbol_id_ = symbol_id;
            CreateSymbol(param.name, SymbolKind::Variable, node, type_hint);
            current_parent_symbol_id_ = param_prev;
        }

        if (node.body)
            node.body->Accept(*this);

        current_parent_symbol_id_ = prev_parent;
        current_function_id_ = prev_function;
        ExitScope();
    }

    void Builder::VisitClassMember(ast::ClassMember& node)
    {
        std::visit([this](auto& member_ptr) {
            if (member_ptr)
                member_ptr->Accept(*this);
        },
                   node.member);
    }

    void Builder::VisitVarStatement(ast::VarStatement& node)
    {
        for (const auto& decl : node.declarations)
        {
            if (decl)
                decl->Accept(*this);
        }
    }

    void Builder::VisitStaticStatement(ast::StaticStatement& node)
    {
        for (const auto& decl : node.declarations)
        {
            if (decl)
                decl->Accept(*this);
        }
    }

    void Builder::VisitGlobalStatement(ast::GlobalStatement& node)
    {
        for (const auto& decl : node.declarations)
        {
            if (decl)
                decl->Accept(*this);
        }
    }

    void Builder::VisitConstStatement(ast::ConstStatement& node)
    {
        std::optional<std::string> type_hint;
        if (node.type_name)
            type_hint = *node.type_name;
        CreateSymbol(node.name, SymbolKind::Constant, node, type_hint);
    }

    void Builder::VisitVarDeclaration(ast::VarDeclaration& node)
    {
        std::optional<std::string> type_hint;
        if (node.type)
            type_hint = node.type->name;
        CreateSymbol(node.name, SymbolKind::Variable, node, type_hint);
    }

    void Builder::VisitStaticDeclaration(ast::StaticDeclaration& node)
    {
        std::optional<std::string> type_hint;
        if (node.type)
            type_hint = node.type->name;
        CreateSymbol(node.name, SymbolKind::Variable, node, type_hint);
    }

    void Builder::VisitGlobalDeclaration(ast::GlobalDeclaration& node)
    {
        std::optional<std::string> type_hint;
        if (node.type)
            type_hint = node.type->name;
        CreateSymbol(node.name, SymbolKind::Variable, node, type_hint);
    }

    void Builder::VisitFieldDeclaration(ast::FieldDeclaration& node)
    {
        std::optional<std::string> type_hint;
        if (node.type)
            type_hint = node.type->name;
        CreateSymbol(node.name, SymbolKind::Field, node, type_hint);
    }

    void Builder::VisitBlockStatement(ast::BlockStatement& node)
    {
        // 为块创建新作用域
        EnterScope(ScopeKind::Block, node.span);
        VisitStatements(node.statements);
        ExitScope();
    }

    void Builder::VisitIfStatement(ast::IfStatement& node)
    {
        for (auto& branch : node.branches)
        {
            if (branch.condition)
                VisitExpression(*branch.condition);
            if (branch.body)
                branch.body->Accept(*this);
        }
    }

    void Builder::VisitForInStatement(ast::ForInStatement& node)
    {
        // 为循环创建作用域
        EnterScope(ScopeKind::Block, node.span);

        CreateSymbol(node.key, SymbolKind::Variable, node);
        CreateSymbol(node.value, SymbolKind::Variable, node);

        if (node.collection)
            VisitExpression(*node.collection);

        if (node.body)
            node.body->Accept(*this);

        ExitScope();
    }

    void Builder::VisitForToStatement(ast::ForToStatement& node)
    {
        // 为循环创建作用域
        EnterScope(ScopeKind::Block, node.span);

        CreateSymbol(node.counter, SymbolKind::Variable, node);

        if (node.start)
            VisitExpression(*node.start);
        if (node.end)
            VisitExpression(*node.end);

        if (node.body)
            node.body->Accept(*this);

        ExitScope();
    }

    void Builder::VisitWhileStatement(ast::WhileStatement& node)
    {
        if (node.condition)
            VisitExpression(*node.condition);

        if (node.body)
            node.body->Accept(*this);
    }

    void Builder::VisitRepeatStatement(ast::RepeatStatement& node)
    {
        EnterScope(ScopeKind::Block, node.span);
        VisitStatements(node.body);

        if (node.condition)
            VisitExpression(*node.condition);

        ExitScope();
    }

    void Builder::VisitCaseStatement(ast::CaseStatement& node)
    {
        if (node.discriminant)
            VisitExpression(*node.discriminant);

        for (auto& branch : node.branches)
        {
            for (auto& value : branch.values)
            {
                if (value)
                    VisitExpression(*value);
            }

            if (branch.body)
                branch.body->Accept(*this);
        }

        if (node.default_case)
            node.default_case->Accept(*this);
    }

    void Builder::VisitTryStatement(ast::TryStatement& node)
    {
        if (node.try_body)
            node.try_body->Accept(*this);

        if (node.except_body)
            node.except_body->Accept(*this);
    }

    void Builder::VisitIdentifier(ast::Identifier& node)
    {
        auto symbol_id = table_.GetScopeManager().FindInScopeChain(current_scope_id_, node.name);
        if (symbol_id)
            table_.AddReference(*symbol_id, node.span, false);
    }

    void Builder::VisitCallExpression(ast::CallExpression& node)
    {
        if (node.callee)
        {
            VisitExpression(*node.callee);

            if (auto* ident = dynamic_cast<ast::Identifier*>(node.callee.get()))
            {
                auto callee_id = table_.GetScopeManager().FindInScopeChain(current_scope_id_, ident->name);
                if (callee_id && current_function_id_)
                    table_.GetCallGraph().AddCall(*current_function_id_, *callee_id, node.span);
            }
        }

        for (auto& arg : node.arguments)
        {
            if (arg.value)
                VisitExpression(*arg.value);
        }
    }

    void Builder::VisitAttributeExpression(ast::AttributeExpression& node)
    {
        if (node.object)
            VisitExpression(*node.object);
    }

    void Builder::VisitAssignmentExpression(ast::AssignmentExpression& node)
    {
        ProcessLeftHandSide(node.left, true);

        if (node.right)
            VisitExpression(*node.right);
    }

    void Builder::VisitAssignmentStatement(ast::AssignmentStatement& node)
    {
        ProcessLeftHandSide(node.left, true);

        if (node.right)
            VisitExpression(*node.right);
    }

    void Builder::VisitBinaryExpression(ast::BinaryExpression& node)
    {
        if (node.left)
            VisitExpression(*node.left);
        if (node.right)
            VisitExpression(*node.right);
    }

    void Builder::VisitUnaryExpression(ast::UnaryExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitTernaryExpression(ast::TernaryExpression& node)
    {
        if (node.condition)
            VisitExpression(*node.condition);
        if (node.consequence)
            VisitExpression(*node.consequence);
        if (node.alternative)
            VisitExpression(*node.alternative);
    }

    void Builder::VisitSubscriptExpression(ast::SubscriptExpression& node)
    {
        if (node.value)
            VisitExpression(*node.value);

        for (auto& index : node.indices)
        {
            if (index.start)
                VisitExpression(*index.start);
            if (index.end)
                VisitExpression(*index.end);
            if (index.step)
                VisitExpression(*index.step);
        }
    }

    void Builder::VisitArrayExpression(ast::ArrayExpression& node)
    {
        for (auto& elem : node.elements)
        {
            if (elem.key)
                VisitExpression(*elem.key);
            if (elem.value)
                VisitExpression(*elem.value);
        }
    }

    void Builder::VisitAnonymousFunctionExpression(ast::AnonymousFunctionExpression& node)
    {
        // 为匿名函数创建作用域
        EnterScope(ScopeKind::Function, node.span);

        for (const auto& param : node.signature.parameters)
        {
            std::optional<std::string> type_hint;
            if (param.type)
                type_hint = param.type->name;
            CreateSymbol(param.name, SymbolKind::Variable, node, type_hint);
        }

        if (node.body)
            node.body->Accept(*this);

        ExitScope();
    }

    void Builder::VisitPrefixIncrementExpression(ast::PrefixIncrementExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitPrefixDecrementExpression(ast::PrefixDecrementExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitPostfixIncrementExpression(ast::PostfixIncrementExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitPostfixDecrementExpression(ast::PostfixDecrementExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitComparisonExpression(ast::ComparisonExpression& node)
    {
        if (node.left)
            VisitExpression(*node.left);

        for (auto& comp : node.comparisons)
        {
            if (comp.right)
                VisitExpression(*comp.right);
        }
    }

    void Builder::VisitExpressionStatement(ast::ExpressionStatement& node)
    {
        if (node.expression)
            VisitExpression(*node.expression);
    }

    void Builder::VisitReturnStatement(ast::ReturnStatement& node)
    {
        if (node.value)
            VisitExpression(*node.value);
    }

    void Builder::VisitUsesStatement(ast::UsesStatement& node)
    {
        // 为每个导入创建一个符号
        for (const auto& unit_name : node.units)
        {
            // 记录导入关系
            table_.AddUnitImport(unit_name, node.span);

            // 创建一个符号来表示这个导入
            [[maybe_unused]] SymbolId import_symbol_id = CreateSymbol(
                unit_name,
                SymbolKind::Module,
                node,
                std::nullopt);
        }
    }

    void Builder::VisitMatrixIterationStatement(ast::MatrixIterationStatement& node)
    {
        // 为矩阵迭代创建作用域
        EnterScope(ScopeKind::Block, node.span);

        if (node.target)
            VisitExpression(*node.target);

        if (node.body)
            node.body->Accept(*this);

        ExitScope();
    }

    void Builder::VisitCompilerDirectiveStatement(ast::CompilerDirectiveStatement& node)
    {
        // 当前阶段：不处理
        // 未来可以在这里维护宏定义状态
    }

    void Builder::VisitConditionalDirectiveStatement(ast::ConditionalDirectiveStatement& node)
    {
        // 当前阶段：不处理
        // 条件指令本身不影响符号收集
    }

    void Builder::VisitConditionalBlockStatement(ast::ConditionalBlockStatement& node)
    {
        // 处理所有分支，不管条件是否满足

        // 处理 THEN 分支 (consequence)
        if (!node.consequence.empty())
        {
            VisitStatements(node.consequence);
        }

        // 处理 ELSE 分支 (alternative)
        if (!node.alternative.empty())
        {
            VisitStatements(node.alternative);
        }
    }

}
