#include <scc_hir_layout.h>

int scc_hir_type_align(scc_hir_module_t *mod, scc_hir_type_ref_t ref,
                       const scc_type_abi_t *abi) {
    scc_hir_type_t *type = scc_hir_module_get_type(mod, ref);
    if (!type)
        return 0;

    switch (type->tag) {
    case SCC_HIR_TYPE_void:
    case SCC_HIR_TYPE_FUNC:
    case SCC_HIR_TYPE_unknown:
        return 1; // void/func 不占空间, 对齐为 1

    case SCC_HIR_TYPE_i8:
    case SCC_HIR_TYPE_u8:
        return scc_type_abi_get_type_align(abi, 1) * 8;
    case SCC_HIR_TYPE_i16:
    case SCC_HIR_TYPE_u16:
    case SCC_HIR_TYPE_f16:
        return scc_type_abi_get_type_align(abi, 2) * 8;
    case SCC_HIR_TYPE_i32:
    case SCC_HIR_TYPE_u32:
    case SCC_HIR_TYPE_f32:
        return scc_type_abi_get_type_align(abi, 4) * 8;
    case SCC_HIR_TYPE_i64:
    case SCC_HIR_TYPE_u64:
    case SCC_HIR_TYPE_f64:
        return scc_type_abi_get_type_align(abi, 8) * 8;
    case SCC_HIR_TYPE_i128:
    case SCC_HIR_TYPE_u128:
    case SCC_HIR_TYPE_f128:
        return scc_type_abi_get_type_align(abi, 16) * 8;

    case SCC_HIR_TYPE_PTR:
        return abi->ptr_align * 8;

    case SCC_HIR_TYPE_ARRAY: {
        return scc_hir_type_align(mod, type->data.array.base, abi);
    }

    case SCC_HIR_TYPE_STRUCT: {
        int max_align_bits = 0;
        for (usize i = 0; i < scc_vec_size(type->data.aggregate.fields); i++) {
            scc_hir_type_ref_t fr = scc_vec_at(type->data.aggregate.fields, i);
            int fa_bits = scc_hir_type_align(mod, fr, abi);
            int fs_bits = scc_hir_type_size(mod, fr, abi);
            int fa2_bits = scc_type_abi_get_field_align(abi, fs_bits / 8, fa_bits / 8) * 8;
            if (fa2_bits > max_align_bits)
                max_align_bits = fa2_bits;
        }
        return scc_type_abi_get_aggregate_align(abi, max_align_bits / 8) * 8;
    }

    case SCC_HIR_TYPE_UNION: {
        int max_align_bits = 0;
        for (usize i = 0; i < scc_vec_size(type->data.aggregate.fields); i++) {
            scc_hir_type_ref_t fr = scc_vec_at(type->data.aggregate.fields, i);
            int fs_bits = scc_hir_type_size(mod, fr, abi);
            int fa_bits = scc_hir_type_align(mod, fr, abi);
            (void)fs_bits;
            if (fa_bits > max_align_bits)
                max_align_bits = fa_bits;
        }
        return scc_type_abi_get_aggregate_align(abi, max_align_bits / 8) * 8;
    }

    default:
        return 1;
    }
}

int scc_hir_type_size(scc_hir_module_t *mod, scc_hir_type_ref_t ref,
                      const scc_type_abi_t *abi) {
    scc_hir_type_t *type = scc_hir_module_get_type(mod, ref);
    if (!type)
        return 0;

    switch (type->tag) {
    case SCC_HIR_TYPE_void:
    case SCC_HIR_TYPE_FUNC:
    case SCC_HIR_TYPE_unknown:
        return 0;

    case SCC_HIR_TYPE_i8:
    case SCC_HIR_TYPE_u8:
        return 8;
    case SCC_HIR_TYPE_i16:
    case SCC_HIR_TYPE_u16:
    case SCC_HIR_TYPE_f16:
        return 16;
    case SCC_HIR_TYPE_i32:
    case SCC_HIR_TYPE_u32:
    case SCC_HIR_TYPE_f32:
        return 32;
    case SCC_HIR_TYPE_i64:
    case SCC_HIR_TYPE_u64:
    case SCC_HIR_TYPE_f64:
        return 64;
    case SCC_HIR_TYPE_i128:
    case SCC_HIR_TYPE_u128:
    case SCC_HIR_TYPE_f128:
        return 128;

    case SCC_HIR_TYPE_PTR:
        return abi->ptr_size * 8;

    case SCC_HIR_TYPE_ARRAY: {
        int es_bits = scc_hir_type_size(mod, type->data.array.base, abi);
        return es_bits * (int)type->data.array.len;
    }

    case SCC_HIR_TYPE_STRUCT: {
        int offset_bits = 0, max_align_bits = 0;
        for (usize i = 0; i < scc_vec_size(type->data.aggregate.fields); i++) {
            scc_hir_type_ref_t fr = scc_vec_at(type->data.aggregate.fields, i);
            int fs_bits = scc_hir_type_size(mod, fr, abi);
            int fa_bits = scc_hir_type_align(mod, fr, abi);
            int fa2_bits = scc_type_abi_get_field_align(abi, fs_bits / 8, fa_bits / 8) * 8;
            offset_bits = SCC_ALIGN_UP(offset_bits, fa2_bits);
            offset_bits += fs_bits;
            if (fa2_bits > max_align_bits)
                max_align_bits = fa2_bits;
        }
        return SCC_ALIGN_UP(offset_bits,
                            scc_type_abi_get_aggregate_align(abi, max_align_bits / 8) * 8);
    }

    case SCC_HIR_TYPE_UNION: {
        int max_size_bits = 0, max_align_bits = 0;
        for (usize i = 0; i < scc_vec_size(type->data.aggregate.fields); i++) {
            scc_hir_type_ref_t fr = scc_vec_at(type->data.aggregate.fields, i);
            int fs_bits = scc_hir_type_size(mod, fr, abi);
            int fa_bits = scc_hir_type_align(mod, fr, abi);
            if (fs_bits > max_size_bits)
                max_size_bits = fs_bits;
            if (fa_bits > max_align_bits)
                max_align_bits = fa_bits;
        }
        return SCC_ALIGN_UP(max_size_bits,
                            scc_type_abi_get_aggregate_align(abi, max_align_bits / 8) * 8);
    }

    default:
        return 0;
    }
}

int scc_hir_field_offset(scc_hir_module_t *mod, scc_hir_type_ref_t ref,
                         int field_idx, const scc_type_abi_t *abi) {
    scc_hir_type_t *type = scc_hir_module_get_type(mod, ref);
    if (!type || (type->tag != SCC_HIR_TYPE_STRUCT &&
                  type->tag != SCC_HIR_TYPE_UNION))
        return 0;

    if (type->tag == SCC_HIR_TYPE_UNION)
        return 0; // union 所有字段偏移为 0

    int offset_bits = 0;
    int count = (int)scc_vec_size(type->data.aggregate.fields);
    if (field_idx >= count)
        return 0;

    for (int i = 0; i < count; i++) {
        scc_hir_type_ref_t fr = scc_vec_at(type->data.aggregate.fields, (usize)i);
        int fs_bits = scc_hir_type_size(mod, fr, abi);
        int fa_bits = scc_hir_type_align(mod, fr, abi);
        int fa2_bits = scc_type_abi_get_field_align(abi, fs_bits / 8, fa_bits / 8) * 8;
        offset_bits = SCC_ALIGN_UP(offset_bits, fa2_bits);
        if (i == field_idx)
            return offset_bits;
        offset_bits += fs_bits;
    }
    return 0;
}

scc_hir_aggregate_layout_t *
scc_hir_aggregate_layout(scc_hir_module_t *mod, scc_hir_type_ref_t ref,
                         const scc_type_abi_t *abi) {
    scc_hir_type_t *type = scc_hir_module_get_type(mod, ref);
    if (!type || (type->tag != SCC_HIR_TYPE_STRUCT &&
                  type->tag != SCC_HIR_TYPE_UNION))
        return NULL;

    int fc = (int)scc_vec_size(type->data.aggregate.fields);

    if (type->tag == SCC_HIR_TYPE_UNION) {
        int max_size_bits = 0, max_align_bits = 0;
        for (usize i = 0; i < scc_vec_size(type->data.aggregate.fields); i++) {
            scc_hir_type_ref_t fr = scc_vec_at(type->data.aggregate.fields, i);
            int fs_bits = scc_hir_type_size(mod, fr, abi);
            int fa_bits = scc_hir_type_align(mod, fr, abi);
            if (fs_bits > max_size_bits)
                max_size_bits = fs_bits;
            if (fa_bits > max_align_bits)
                max_align_bits = fa_bits;
        }

        scc_hir_aggregate_layout_t *al =
            scc_malloc(sizeof(scc_hir_aggregate_layout_t) +
                       (usize)fc * sizeof(scc_hir_field_layout_t));
        al->size = SCC_ALIGN_UP(max_size_bits, max_align_bits);
        al->align = max_align_bits;
        al->field_count = fc;
        for (int i = 0; i < fc; i++) {
            scc_hir_type_ref_t fr =
                scc_vec_at(type->data.aggregate.fields, (usize)i);
            int fs_bits = scc_hir_type_size(mod, fr, abi);
            int fa_bits = scc_hir_type_align(mod, fr, abi);
            al->fields[i].offset = 0;
            al->fields[i].size = fs_bits;
            al->fields[i].align = fa_bits;
        }
        return al;
    }

    // STRUCT
    scc_hir_aggregate_layout_t *al =
        scc_malloc(sizeof(scc_hir_aggregate_layout_t) +
                   (usize)fc * sizeof(scc_hir_field_layout_t));
    al->field_count = fc;

    int offset_bits = 0, max_align_bits = 0;
    for (int i = 0; i < fc; i++) {
        scc_hir_type_ref_t fr =
            scc_vec_at(type->data.aggregate.fields, (usize)i);
        int fs_bits = scc_hir_type_size(mod, fr, abi);
        int fa_bits = scc_hir_type_align(mod, fr, abi);
        int fa2_bits = scc_type_abi_get_field_align(abi, fs_bits / 8, fa_bits / 8) * 8;
        offset_bits = SCC_ALIGN_UP(offset_bits, fa2_bits);

        al->fields[i].offset = offset_bits;
        al->fields[i].size = fs_bits;
        al->fields[i].align = fa_bits;

        offset_bits += fs_bits;
        if (fa2_bits > max_align_bits)
            max_align_bits = fa2_bits;
    }

    al->size = SCC_ALIGN_UP(offset_bits,
                            scc_type_abi_get_aggregate_align(abi, max_align_bits / 8) * 8);
    al->align = scc_type_abi_get_aggregate_align(abi, max_align_bits / 8) * 8;
    return al;
}

void scc_hir_aggregate_layout_free(scc_hir_aggregate_layout_t *layout) {
    scc_free(layout);
}