#include "AddressCalculationExpansion.h"
#include <iostream>
#include <vector>
#include "IR.h"
#include "IRBuilder.h"

extern int DEBUG;

namespace sysy {

bool AddressCalculationExpansion::run() {
    bool changed = false;

    for (auto& funcPair : pModule->getFunctions()) {
        Function* func = funcPair.second.get();
        for (auto& bb_ptr : func->getBasicBlocks()) {
            BasicBlock* bb = bb_ptr.get();
            for (auto it = bb->getInstructions().begin(); it != bb->getInstructions().end(); ) {
                Instruction* inst = it->get();

                Value* basePointer = nullptr;
                Value* valueToStore = nullptr;
                size_t firstIndexOperandIdx = 0;
                size_t numBaseOperands = 0;

                if (inst->isLoad()) {
                    numBaseOperands = 1;
                    basePointer = inst->getOperand(0);
                    firstIndexOperandIdx = 1;
                } else if (inst->isStore()) {
                    numBaseOperands = 2;
                    valueToStore = inst->getOperand(0);
                    basePointer = inst->getOperand(1);
                    firstIndexOperandIdx = 2;
                } else {
                    ++it;
                    continue;
                }

                if (inst->getNumOperands() <= numBaseOperands) {
                    ++it;
                    continue;
                }

                std::vector<int> dims;
                if (AllocaInst* allocaInst = dynamic_cast<AllocaInst*>(basePointer)) {
                    for (const auto& use_ptr : allocaInst->getDims()) {
                        Value* dimValue = use_ptr->getValue();
                        if (ConstantValue* constVal = dynamic_cast<ConstantValue*>(dimValue)) {
                            dims.push_back(constVal->getInt());
                        } else {
                            std::cerr << "Warning: AllocaInst dimension is not a constant integer. Skipping GEP expansion for: ";
                            SysYPrinter::printValue(allocaInst);
                            std::cerr << "\n";
                            dims.clear();
                            break;
                        }
                    }
                } else if (GlobalValue* globalValue = dynamic_cast<GlobalValue*>(basePointer)) {
                    // 遍历 GlobalValue 的所有维度操作数
                    for (const auto& use_ptr : globalValue->getDims()) {
                        Value* dimValue = use_ptr->getValue();
                        // 将维度值转换为常量整数
                        if (ConstantInteger* constVal = dynamic_cast<ConstantInteger*>(dimValue)) {
                            dims.push_back(constVal->getInt());
                        } else {
                            // 如果维度不是常量整数，则无法处理。
                            // 根据 IR.h 中 GlobalValue 的构造函数，这种情况不应发生，但作为安全检查是好的。
                            std::cerr << "Warning: GlobalValue dimension is not a constant integer. Skipping GEP expansion for: ";
                            SysYPrinter::printValue(globalValue);
                            std::cerr << "\n";
                            dims.clear(); // 清空已收集的部分维度信息
                            break;
                        }
                    }
                } else {
                    std::cerr << "Warning: Base pointer is not AllocaInst/GlobalValue or its array dimensions cannot be determined for GEP expansion. Skipping GEP for: ";
                    SysYPrinter::printValue(basePointer);
                    std::cerr << " in instruction ";
                    SysYPrinter::printInst(inst);
                    std::cerr << "\n";
                    ++it;
                    continue;
                }

                if (dims.empty() && (inst->getNumOperands() > numBaseOperands)) {
                    if (DEBUG) {
                        std::cerr << "ACE Warning: Could not get valid array dimensions for ";
                        SysYPrinter::printValue(basePointer);
                        std::cerr << " in instruction ";
                        SysYPrinter::printInst(inst);
                        std::cerr << " (expected dimensions for indices, but got none).\n";
                    }
                    ++it;
                    continue;
                }

                std::vector<Value*> indexOperands;
                for (size_t i = firstIndexOperandIdx; i < inst->getNumOperands(); ++i) {
                    indexOperands.push_back(inst->getOperand(i));
                }

                if (AllocaInst* allocaInst = dynamic_cast<AllocaInst*>(basePointer)) {
                    if (allocaInst->getNumDims() != indexOperands.size()) {
                        if (DEBUG) {
                            std::cerr << "ACE Warning: Index count (" << indexOperands.size() << ") does not match AllocaInst dimensions (" << allocaInst->getNumDims() << ") for instruction ";
                            SysYPrinter::printInst(inst);
                            std::cerr << "\n";
                        }
                        ++it;
                        continue;
                    }
                }

                Value* totalOffset = ConstantInteger::get(0);
                pBuilder->setPosition(bb, it);

                for (size_t i = 0; i < indexOperands.size(); ++i) {
                    Value* index = indexOperands[i];
                    int stride = calculateStride(dims, i);
                    Value* strideConst = ConstantInteger::get(stride);
                    Type* intType = Type::getIntType();
                    BinaryInst* currentDimOffsetInst = pBuilder->createBinaryInst(Instruction::kMul, intType, index, strideConst);
                    BinaryInst* newTotalOffsetInst = pBuilder->createBinaryInst(Instruction::kAdd, intType, totalOffset, currentDimOffsetInst);
                    totalOffset = newTotalOffsetInst;
                }

                // 计算有效地址：effective_address = basePointer + totalOffset
                Value* effective_address = pBuilder->createBinaryInst(Instruction::kAdd, basePointer->getType(), basePointer, totalOffset);

                // 创建新的 LoadInst 或 StoreInst，indices 为空
                Instruction* newInst = nullptr;
                if (inst->isLoad()) {
                    newInst = pBuilder->createLoadInst(effective_address, {});
                    inst->replaceAllUsesWith(newInst);
                } else { // StoreInst
                    newInst = pBuilder->createStoreInst(valueToStore, effective_address, {});
                }

                Instruction* oldInst = it->get();
                ++it;

                for (size_t i = 0; i < oldInst->getNumOperands(); ++i) {
                    Value* operandValue = oldInst->getOperand(i);
                    if (operandValue) {
                        for (auto use_it = operandValue->getUses().begin(); use_it != operandValue->getUses().end(); ++use_it) {
                            if ((*use_it)->getUser() == oldInst && (*use_it)->getIndex() == i) {
                                operandValue->removeUse(*use_it);
                                break;
                            }
                        }
                    }
                }

                bb->getInstructions().erase(std::prev(it));
                changed = true;

                if (DEBUG) {
                    std::cerr << "ACE: Computed effective address:\n";
                    SysYPrinter::printInst(dynamic_cast<Instruction*>(effective_address));
                    std::cerr << "ACE: New Load/Store instruction:\n";
                    SysYPrinter::printInst(newInst);
                    std::cerr << "--------------------------------\n";
                }
            }
        }
    }
    return changed;
}

} // namespace sysy