// 当前仅支持 void、i32 和 i32*。
#include "ir/IR.h"
#include <cassert>

namespace ir {

Type::Type(Kind k) : kind_(k) {}

size_t Type::Size() const {
    switch (kind_) {
        case Kind::Void:    return 0;
        case Kind::Int32:   return 4;
        case Kind::Float:   return 4;          // 单精度浮点 4 字节
        case Kind::PtrInt32: return 8;         // 假设 64 位指针
        case Kind::PtrFloat: return 8;
        case Kind::Label:   return 8;          // 标签地址大小（指针大小）
        case Kind::PtrInt1: return 8;         // 指向 i1 的指针大小
        default:            return 0;          // 派生类应重写
    }
}

size_t Type::Alignment() const {
    switch (kind_) {
        case Kind::Int32:   return 4;
        case Kind::Float:   return 4;
        case Kind::PtrInt32: return 8;
        case Kind::PtrFloat: return 8;
        case Kind::Label:   return 8;          // 与指针相同
        default:            return 1;
    }
}

bool Type::IsComplete() const {
    return kind_ != Kind::Void;
}
const std::shared_ptr<Type>& Type::GetVoidType() {
  static const std::shared_ptr<Type> type = std::shared_ptr<Type>(new Type(Kind::Void));
  return type;
}

const std::shared_ptr<Type>& Type::GetInt32Type() {
  static const std::shared_ptr<Type> type = std::shared_ptr<Type>(new Type(Kind::Int32));
  return type;
}

const std::shared_ptr<Type>& Type::GetFloatType() {
    static const std::shared_ptr<Type> type(new Type(Kind::Float));
    return type;
}

const std::shared_ptr<Type>& Type::GetPtrInt32Type() {
  static const std::shared_ptr<Type> type = std::shared_ptr<Type>(new Type(Kind::PtrInt32));
  return type;
}

const std::shared_ptr<Type>& Type::GetPtrFloatType() {
    static const std::shared_ptr<Type> type(new Type(Kind::PtrFloat));
    return type;
}

const std::shared_ptr<Type>& Type::GetLabelType() {
    static const std::shared_ptr<Type> type(new Type(Kind::Label));
    return type;
}

// Int1 类型表示布尔值，通常用于比较指令的结果
const std::shared_ptr<Type>& Type::GetInt1Type() {
    static const std::shared_ptr<Type> type = std::shared_ptr<Type>(new Type(Kind::Int1));
    return type;
  }
// PtrInt1 类型表示指向 Int1 的指针，主要用于条件跳转等场景
const std::shared_ptr<Type>& Type::GetPtrInt1Type() {
    static const std::shared_ptr<Type> type = std::shared_ptr<Type>(new Type(Kind::PtrInt1));
    return type;
  }

// ---------- 数组类型缓存 ----------
// 使用自定义键类型保证唯一性：元素类型指针 + 维度向量
struct ArrayKey {
    const Type* elem_type;
    std::vector<int> dims;

    bool operator==(const ArrayKey& other) const {
        return elem_type == other.elem_type && dims == other.dims;
    }
};

struct ArrayKeyHash {
    std::size_t operator()(const ArrayKey& key) const {
        std::size_t h = std::hash<const Type*>{}(key.elem_type);
        for (int d : key.dims) {
            h ^= std::hash<int>{}(d) + 0x9e3779b9 + (h << 6) + (h >> 2);
        }
        return h;
    }
};

static std::unordered_map<ArrayKey, std::weak_ptr<ArrayType>, ArrayKeyHash>& GetArrayCache() {
    static std::unordered_map<ArrayKey, std::weak_ptr<ArrayType>, ArrayKeyHash> cache;
    return cache;
}

std::shared_ptr<ArrayType> Type::GetArrayType(std::shared_ptr<Type> elem,
                                              std::vector<int> dims) {
    // 检查维度合法性
    for (int d : dims) {
        if (d <= 0) {
            // SysY 数组维度必须为正整数常量表达式，这里假设已检查
            return nullptr;
        }
    }

    ArrayKey key{elem.get(), dims};
    auto& cache = GetArrayCache();
    auto it = cache.find(key);
    if (it != cache.end()) {
        auto ptr = it->second.lock();
        if (ptr) return ptr;
    }

    auto arr = std::shared_ptr<ArrayType>(new ArrayType(std::move(elem), std::move(dims)));
    cache[key] = arr;
    return arr;
}

// ---------- 函数类型缓存 ----------
struct FunctionKey {
    const Type* return_type;
    std::vector<const Type*> param_types;

    bool operator==(const FunctionKey& other) const {
        return return_type == other.return_type && param_types == other.param_types;
    }
};

struct FunctionKeyHash {
    std::size_t operator()(const FunctionKey& key) const {
        std::size_t h = std::hash<const Type*>{}(key.return_type);
        for (const Type* t : key.param_types) {
            h ^= std::hash<const Type*>{}(t) + 0x9e3779b9 + (h << 6) + (h >> 2);
        }
        return h;
    }
};

static std::unordered_map<FunctionKey, std::weak_ptr<FunctionType>, FunctionKeyHash>& GetFunctionCache() {
    static std::unordered_map<FunctionKey, std::weak_ptr<FunctionType>, FunctionKeyHash> cache;
    return cache;
}

std::shared_ptr<FunctionType> Type::GetFunctionType(std::shared_ptr<Type> ret,
                                                    std::vector<std::shared_ptr<Type>> params) {
    // 提取裸指针用于键（保证唯一性，因为 shared_ptr 指向同一对象）
    std::vector<const Type*> param_ptrs;
    param_ptrs.reserve(params.size());
    for (const auto& p : params) {
        param_ptrs.push_back(p.get());
    }

    FunctionKey key{ret.get(), std::move(param_ptrs)};
    auto& cache = GetFunctionCache();
    auto it = cache.find(key);
    if (it != cache.end()) {
        auto ptr = it->second.lock();
        if (ptr) return ptr;
    }

    auto func = std::shared_ptr<FunctionType>(new FunctionType(std::move(ret), std::move(params)));
    cache[key] = func;
    return func;
}

// ---------- ArrayType 实现 ----------
ArrayType::ArrayType(std::shared_ptr<Type> elem, std::vector<int> dims)
    : Type(Kind::Array), elem_type_(std::move(elem)), dims_(std::move(dims)) {
    // 数组元素类型必须是完整类型
    assert(elem_type_ && elem_type_->IsComplete());
}

size_t ArrayType::GetElementCount() const {
    size_t count = 1;
    for (int d : dims_) count *= d;
    return count;
}

size_t ArrayType::Size() const {
    return GetElementCount() * elem_type_->Size();
}

size_t ArrayType::Alignment() const {
    // 数组对齐等于其元素对齐
    return elem_type_->Alignment();
}

bool ArrayType::IsComplete() const {
    // 维度已确定且元素类型完整，则数组完整
    return !dims_.empty() && elem_type_->IsComplete();
}

// ---------- FunctionType 实现 ----------
FunctionType::FunctionType(std::shared_ptr<Type> ret,
                           std::vector<std::shared_ptr<Type>> params)
    : Type(Kind::Function), return_type_(std::move(ret)), param_types_(std::move(params)) {}

size_t FunctionType::Size() const {
    // 函数类型没有运行时大小，通常用于类型检查，返回 0
    return 0;
}

size_t FunctionType::Alignment() const {
    // 不对齐
    return 1;
}

bool FunctionType::IsComplete() const {
    // 函数类型总是完整的（只要返回类型完整，但 SysY 中 void 也视为完整）
    return true;
}

}  // namespace ir