#include "spl_layout.h"

static void infos_grow(spl_layout_t *layout, usize n) {
    while (vec_size(layout->infos) < n) {
        spl_layout_info_t info = {0};
        vec_init(info.field_bit_offsets);
        vec_init(info.field_bit_sizes);
        vec_push(layout->infos, info);
    }
}

static spl_layout_info_t *infos_at(spl_layout_t *layout, spl_type_id_t id) {
    if (id == 0 || id >= vec_size(layout->infos))
        return NULL;
    return &vec_at(layout->infos, id);
}

void spl_layout_add_scalar(spl_layout_t *layout, usize size, usize abi_align) {
    spl_layout_scalar_t s = {0};
    s.size = size;
    s.abi_align = abi_align;
    s.pref_align = abi_align;
    vec_init(s.invalid_values);
    vec_push(layout->scalars, s);
}

void spl_layout_add_ptr(spl_layout_t *layout, usize addr_space, usize size, usize align) {
    spl_layout_ptr_t p = {0};
    p.add_space = addr_space;
    p.ptr_size = size;
    p.ptr_align = align;
    vec_push(layout->pointers, p);
}

void spl_layout_init(spl_layout_t *layout, spl_layout_target_t *target) {
    vec_init(layout->scalars);
    vec_init(layout->pointers);
    vec_init(layout->infos);
    layout->target = target;

    /* 默认聚合策略：按最大字段对齐，不允许重排 */
    layout->aggregate.allow_field_reordering = false;
    layout->aggregate.align_rule = ALIGN_MAX;
    layout->aggregate.fixed_align = 0;

    /* 常见 64 位平台的默认布局。
       若目标平台不适用，调用者 clear 后自行 add_*。 */
    spl_layout_add_scalar(layout, 1, 1); /* i8 / u8 / bool */
    spl_layout_add_scalar(layout, 2, 2); /* i16 / u16 */
    spl_layout_add_scalar(layout, 4, 4); /* i32 / u32 / f32 */
    spl_layout_add_scalar(layout, 8, 8); /* i64 / u64 / f64 */

    usize ps = target && target->ptr_size ? target->ptr_size : sizeof(void *);
    usize pa = target && target->ptr_align ? target->ptr_align : sizeof(void *);
    spl_layout_add_ptr(layout, 0, ps, pa);
}

void spl_layout_drop(spl_layout_t *layout) {
    for (usize i = 0; i < vec_size(layout->scalars); ++i)
        vec_free(vec_at(layout->scalars, i).invalid_values);
    vec_free(layout->scalars);

    vec_free(layout->pointers);

    for (usize i = 0; i < vec_size(layout->infos); ++i) {
        vec_free(vec_at(layout->infos, i).field_bit_offsets);
        vec_free(vec_at(layout->infos, i).field_bit_sizes);
    }
    vec_free(layout->infos);
}

static const spl_layout_scalar_t *scalar_for_size(spl_layout_t *layout, usize size_bytes) {
    for (usize i = 0; i < vec_size(layout->scalars); ++i) {
        spl_layout_scalar_t *s = &vec_at(layout->scalars, i);
        if (s->size == size_bytes)
            return s;
    }
    return NULL;
}

static const spl_layout_ptr_t *default_ptr(spl_layout_t *layout) {
    if (vec_size(layout->pointers) == 0)
        return NULL;
    return &vec_at(layout->pointers, 0);
}

static bool layout_scalar(spl_layout_t *layout, spl_type_node_t *tn, spl_layout_info_t *out) {
    usize size_bits = 0;
    switch (tn->kind) {
    case SPL_TYPE_INT:
        if (tn->int_type.bits == 0)
            return false; /* comptime_int 未落地 */
        size_bits = tn->int_type.bits;
        break;
    case SPL_TYPE_FLOAT:
        if (tn->float_type.bits == 0)
            return false; /* comptime_float 未落地 */
        size_bits = tn->float_type.bits;
        break;
    case SPL_TYPE_BOOL:
        size_bits = 8;
        break;
    default:
        return false;
    }
    usize size_bytes = (size_bits + 7) / 8;
    const spl_layout_scalar_t *sl = scalar_for_size(layout, size_bytes);
    if (!sl)
        return false;

    out->size_bits = sl->size * 8;
    out->align_bits = sl->abi_align * 8;
    return true;
}

static bool layout_ptr(spl_layout_t *layout, spl_layout_info_t *out) {
    const spl_layout_ptr_t *pl = default_ptr(layout);
    if (!pl)
        return false;
    out->size_bits = pl->ptr_size * 8;
    out->align_bits = pl->ptr_align * 8;
    return true;
}

static usize round_up_bits(usize x, usize align_bits) {
    if (align_bits <= 1)
        return x;
    usize a = align_bits;
    return (x + a - 1) / a * a;
}

static bool layout_agg(spl_layout_t *layout, spl_type_node_t *tn, spl_layout_info_t *out) {
    usize n = vec_size(tn->agg_type.variants);

    if (tn->kind == SPL_TYPE_ENUM) {
        spl_type_id_t tag_tid = tn->agg_type.tag_type;
        if (!tag_tid)
            return false;
        spl_layout_info_t *tag_info = infos_at(layout, tag_tid);
        if (!tag_info || !tag_info->resolved)
            return false;

        usize tag_align_bits = tag_info->align_bits ? tag_info->align_bits : 8;
        usize payload_align_bits = 8;
        usize payload_size_bits = 0;

        for (usize i = 0; i < n; ++i) {
            spl_type_member_t *m = &vec_at(tn->agg_type.variants, i);
            if (m->type_id == 0) /* 无载荷变体 */
                continue;
            spl_layout_info_t *fi = infos_at(layout, m->type_id);
            if (!fi || !fi->resolved)
                return false;
            usize fa = fi->align_bits ? fi->align_bits : 8;
            if (fa > payload_align_bits)
                payload_align_bits = fa;
            if (fi->size_bits > payload_size_bits)
                payload_size_bits = fi->size_bits;
        }

        usize payload_off = round_up_bits(tag_info->size_bits, payload_align_bits);
        usize final_align_bits =
            tag_align_bits > payload_align_bits ? tag_align_bits : payload_align_bits;
        out->size_bits = round_up_bits(payload_off + payload_size_bits, final_align_bits);
        out->align_bits = final_align_bits;

        /* 字段索引约定（与 ir2vm 一致）：0 = tag，1..n = 各变体载荷（同一偏移）。 */
        vec_push(out->field_bit_offsets, 0);
        vec_push(out->field_bit_sizes, tag_info->size_bits);
        for (usize i = 0; i < n; ++i) {
            spl_type_member_t *m = &vec_at(tn->agg_type.variants, i);
            usize fsize = 0;
            if (m->type_id) {
                spl_layout_info_t *fi = infos_at(layout, m->type_id);
                if (fi)
                    fsize = fi->size_bits;
            }
            vec_push(out->field_bit_offsets, payload_off);
            vec_push(out->field_bit_sizes, fsize);
        }
        return true;
    }

    if (n == 0) {
        out->size_bits = 0;
        out->align_bits = 8;
        return true;
    }

    usize max_align_bits = 8;
    usize offset_bits = 0;
    bool is_union = (tn->kind == SPL_TYPE_UNION);

    for (usize i = 0; i < n; ++i) {
        spl_type_member_t *m = &vec_at(tn->agg_type.variants, i);
        usize fsize_bits = 0;
        usize fsize_align = 8;
        if (m->type_id) {
            spl_layout_info_t *fi = infos_at(layout, m->type_id);
            if (!fi || !fi->resolved)
                return false;
            fsize_bits = fi->size_bits;
            fsize_align = fi->align_bits ? fi->align_bits : 8;
        }
        usize field_align_bits = fsize_align;

        if (is_union) {
            vec_push(out->field_bit_offsets, 0);
            vec_push(out->field_bit_sizes, fsize_bits);
            if (fsize_bits > out->size_bits)
                out->size_bits = fsize_bits;
        } else {
            offset_bits = round_up_bits(offset_bits, field_align_bits);
            vec_push(out->field_bit_offsets, offset_bits);
            vec_push(out->field_bit_sizes, fsize_bits);
            offset_bits += fsize_bits;
        }
        if (field_align_bits > max_align_bits)
            max_align_bits = field_align_bits;
    }

    usize final_align_bits = max_align_bits;
    if (layout->aggregate.align_rule == ALIGN_FIXED) {
        final_align_bits = layout->aggregate.fixed_align * 8;
    } else if (layout->aggregate.align_rule == ALIGN_MAX_WITH_MIN) {
        usize min_bits = layout->aggregate.fixed_align * 8;
        if (final_align_bits < min_bits)
            final_align_bits = min_bits;
    }

    if (is_union)
        out->size_bits = round_up_bits(out->size_bits, final_align_bits);
    else
        out->size_bits = round_up_bits(offset_bits, final_align_bits);
    out->align_bits = final_align_bits;
    return true;
}

static void resolve_type(spl_layout_t *layout, spl_type_t *type, spl_type_id_t id) {
    spl_layout_info_t *info = infos_at(layout, id);
    if (!info || info->resolved)
        return;

    spl_type_node_t *tn = spl_type_node(type, id);
    if (!tn)
        return;

    /* 先标记已解析：自引用类型（通过 ptr）不会无限递归，
       但保持防御性写法。 */
    info->resolved = true;

    switch (tn->kind) {
    case SPL_TYPE_ERROR:
    case SPL_TYPE_INFER:
    case SPL_TYPE_FORALL:
    case SPL_TYPE_EXISTS:
    case SPL_TYPE_NEVER:
    case SPL_TYPE_UNIT:
        info->size_bits = 0;
        info->align_bits = 8;
        return;

    case SPL_TYPE_ANY:
        /* 顶层类型：stage1 按不透明指针处理 */
        layout_ptr(layout, info);
        return;

    case SPL_TYPE_BOOL:
    case SPL_TYPE_INT:
    case SPL_TYPE_FLOAT:
        layout_scalar(layout, tn, info);
        return;

    case SPL_TYPE_PTR:
        layout_ptr(layout, info);
        return;

    case SPL_TYPE_ARRAY: {
        resolve_type(layout, type, tn->array_type.element);
        spl_layout_info_t *elem = infos_at(layout, tn->array_type.element);
        if (!elem || !elem->resolved)
            return;
        info->size_bits = elem->size_bits * tn->array_type.len;
        info->align_bits = elem->align_bits;
        return;
    }

    case SPL_TYPE_STRUCT:
    case SPL_TYPE_UNION:
    case SPL_TYPE_ENUM:
    case SPL_TYPE_SHAPE: {
        if (tn->kind == SPL_TYPE_ENUM && tn->agg_type.tag_type)
            resolve_type(layout, type, tn->agg_type.tag_type);
        for (usize i = 0; i < vec_size(tn->agg_type.variants); ++i) {
            spl_type_member_t *m = &vec_at(tn->agg_type.variants, i);
            if (m->type_id)
                resolve_type(layout, type, m->type_id);
        }
        layout_agg(layout, tn, info);
        return;
    }

    case SPL_TYPE_FN:
        /* 函数类型按函数指针布局 */
        for (usize i = 0; i < vec_size(tn->fn_type.params); ++i)
            resolve_type(layout, type, vec_at(tn->fn_type.params, i));
        resolve_type(layout, type, tn->fn_type.ret);
        layout_ptr(layout, info);
        return;
    }
}

void spl_layout_resolve(spl_layout_t *layout, spl_type_t *type) {
    infos_grow(layout, vec_size(type->type_table));
    for (usize i = 1; i < vec_size(type->type_table); ++i)
        resolve_type(layout, type, i);
}

const spl_layout_info_t *spl_layout_info_of(spl_layout_t *layout, spl_type_id_t id) {
    spl_layout_info_t *info = infos_at(layout, id);
    if (!info || !info->resolved)
        return NULL;
    return info;
}