// Copyright (c) 2026 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.

module abi

import v2.mir
import v2.pref

pub enum X64Abi {
    sysv
    windows
}

const sysv_int_arg_reg_count = 6
const sysv_sse_arg_reg_count = 8
const sysv_max_direct_aggregate_size = 16

struct SysVLocationState {
mut:
    int_regs    int
    sse_regs    int
    stack_slots int
}

// lower annotates MIR with ABI classification metadata.
// Current scope is intentionally conservative: it classifies which arguments
// and return values must be passed indirectly for each function.
pub fn lower(mut m mir.Module, arch pref.Arch) {
    lower_with_x64_abi(mut m, arch, .sysv)
}

pub fn lower_with_x64_abi(mut m mir.Module, arch pref.Arch, x64_abi X64Abi) {
    is_x64 := arch == .x64
    mut fn_by_name := map[string]int{}
    for i := 0; i < m.funcs.len; i++ {
        mut f := &m.funcs[i]
        fn_by_name[f.name] = i
        if is_x64 {
            f.abi_ret_class = abi_value_class(m, f.typ, arch, x64_abi)
            f.abi_ret_indirect = abi_class_is_indirect(f.abi_ret_class, m, f.typ, arch, x64_abi)
        } else {
            f.abi_ret_indirect = needs_indirect(m, f.typ, arch, x64_abi)
        }
        f.abi_param_class = []mir.AbiArgClass{len: f.params.len, init: .in_reg}
        if is_x64 {
            f.abi_param_classes = []mir.AbiValueClass{len: f.params.len}
            f.abi_param_layouts = []mir.AbiValueLayout{len: f.params.len}
        }
        mut param_loc_state := SysVLocationState{}
        if arch == .x64 && x64_abi == .sysv && f.abi_ret_indirect {
            param_loc_state.int_regs = 1
        }
        for pi, param_id in f.params {
            if param_id < 0 || param_id >= m.values.len {
                continue
            }
            param_typ := m.values[param_id].typ
            if is_x64 {
                param_class := abi_value_class(m, param_typ, arch, x64_abi)
                f.abi_param_classes[pi] = param_class
                if x64_abi == .sysv {
                    f.abi_param_layouts[pi] = sysv_assign_value_layout(param_class, mut
                        param_loc_state)
                }
                if abi_class_is_indirect(param_class, m, param_typ, arch, x64_abi) {
                    f.abi_param_class[pi] = .indirect
                }
            } else if needs_indirect(m, param_typ, arch, x64_abi) {
                f.abi_param_class[pi] = .indirect
            }
        }
    }

    lower_calls(mut m, arch, x64_abi, fn_by_name)
}

fn needs_indirect(m mir.Module, typ_id int, arch pref.Arch, x64_abi X64Abi) bool {
    ssa_mod := m.ssa()
    if ssa_mod == unsafe { nil } || typ_id <= 0 || typ_id >= ssa_mod.type_store.types.len {
        return false
    }
    typ := ssa_mod.type_store.types[typ_id]
    if typ.kind !in [.struct_t, .array_t] {
        return false
    }
    size := m.type_size(typ_id)
    return match arch {
        .arm64 {
            size > 16
        }
        .x64 {
            match x64_abi {
                .sysv { size > sysv_max_direct_aggregate_size }
                .windows { size !in [1, 2, 4, 8] }
            }
        }
        else {
            size > 16
        }
    }
}

fn abi_class_is_indirect(value_class mir.AbiValueClass, m mir.Module, typ_id int, arch pref.Arch, x64_abi X64Abi) bool {
    if value_class.mode == .indirect {
        return true
    }
    return needs_indirect(m, typ_id, arch, x64_abi)
}

fn abi_value_class(m mir.Module, typ_id int, arch pref.Arch, x64_abi X64Abi) mir.AbiValueClass {
    ssa_mod := m.ssa()
    size := m.type_size(typ_id)
    if ssa_mod == unsafe { nil } || typ_id <= 0 || typ_id >= ssa_mod.type_store.types.len {
        return mir.AbiValueClass{
            mode: .direct
            size: size
        }
    }
    typ := ssa_mod.type_store.types[typ_id]
    if arch == .x64 && x64_abi == .sysv {
        if typ.kind !in [.struct_t, .array_t] {
            return sysv_scalar_abi_value_class(m, typ_id)
        }
        return sysv_abi_value_class(m, typ_id)
    }
    indirect := needs_indirect(m, typ_id, arch, x64_abi)
    return mir.AbiValueClass{
        mode: if indirect { .indirect } else { .direct }
        size: size
    }
}

fn sysv_scalar_abi_value_class(m mir.Module, typ_id int) mir.AbiValueClass {
    ssa_mod := m.ssa()
    size := m.type_size(typ_id)
    if ssa_mod == unsafe { nil } || typ_id <= 0 || typ_id >= ssa_mod.type_store.types.len {
        return mir.AbiValueClass{
            mode: .direct
            size: size
        }
    }
    typ := ssa_mod.type_store.types[typ_id]
    classes := match typ.kind {
        .int_t, .ptr_t, .func_t {
            [mir.AbiEightbyteClass.integer]
        }
        .float_t {
            match typ.width {
                32, 64 { [mir.AbiEightbyteClass.sse] }
                128 { [mir.AbiEightbyteClass.sse, .sseup] }
                else { []mir.AbiEightbyteClass{} }
            }
        }
        else {
            []mir.AbiEightbyteClass{}
        }
    }

    return mir.AbiValueClass{
        mode:    .direct
        size:    size
        classes: classes
    }
}

fn sysv_abi_value_class(m mir.Module, typ_id int) mir.AbiValueClass {
    size := m.type_size(typ_id)
    if size <= 0 {
        return mir.AbiValueClass{
            mode: .direct
            size: size
        }
    }
    if sysv_aggregate_must_be_memory_before_classification(size) {
        return sysv_memory_value_class(size)
    }
    mut classes := []mir.AbiEightbyteClass{len: (size + 7) / 8, init: .no_class}
    mut visiting := map[int]bool{}
    if !sysv_classify_type_into(m, typ_id, 0, mut classes, mut visiting) {
        return sysv_memory_value_class(size)
    }
    classes = sysv_post_merge_classes(size, classes)
    if classes.len == 1 && classes[0] == .memory {
        return sysv_memory_value_class(size)
    }
    return mir.AbiValueClass{
        mode:    .direct
        size:    size
        classes: classes
    }
}

fn sysv_aggregate_must_be_memory_before_classification(size int) bool {
    // Current SSA aggregate types do not model a SysV vector aggregate larger
    // than two eightbytes, so larger struct/array aggregates cannot be passed
    // in registers and should not be classified element by element.
    return size > sysv_max_direct_aggregate_size
}

fn sysv_memory_value_class(size int) mir.AbiValueClass {
    return mir.AbiValueClass{
        mode:    .indirect
        size:    size
        classes: [.memory]
    }
}

fn sysv_classify_type_into(m mir.Module, typ_id int, byte_offset int, mut classes []mir.AbiEightbyteClass, mut visiting map[int]bool) bool {
    ssa_mod := m.ssa()
    if ssa_mod == unsafe { nil } || typ_id <= 0 || typ_id >= ssa_mod.type_store.types.len {
        return true
    }
    if visiting[typ_id] {
        return sysv_merge_class_range(mut classes, byte_offset, 8, .integer)
    }
    visiting[typ_id] = true
    typ := ssa_mod.type_store.types[typ_id]
    ok := match typ.kind {
        .int_t, .ptr_t, .func_t {
            sysv_merge_class_range(mut classes, byte_offset, m.type_size(typ_id), .integer)
        }
        .float_t {
            match typ.width {
                32, 64 {
                    sysv_merge_class_range(mut classes, byte_offset, m.type_size(typ_id), .sse)
                }
                80 {
                    false
                }
                128 {
                    sysv_merge_class_range(mut classes, byte_offset, 8, .sse)
                        && sysv_merge_class_range(mut classes, byte_offset + 8, 8, .sseup)
                }
                else {
                    false
                }
            }
        }
        .array_t {
            sysv_classify_array_into(m, typ.elem_type, typ.len, byte_offset, mut classes, mut
                visiting)
        }
        .struct_t {
            sysv_classify_struct_into(m, typ.fields, byte_offset, mut classes, mut visiting)
        }
        .void_t, .label_t, .metadata_t {
            true
        }
    }

    visiting.delete(typ_id)
    return ok
}

fn sysv_classify_array_into(m mir.Module, elem_type int, len int, byte_offset int, mut classes []mir.AbiEightbyteClass, mut visiting map[int]bool) bool {
    elem_size := m.type_size(elem_type)
    for i := 0; i < len; i++ {
        if !sysv_classify_type_into(m, elem_type, byte_offset + i * elem_size, mut classes, mut
            visiting) {
            return false
        }
    }
    return true
}

fn sysv_classify_struct_into(m mir.Module, fields []int, byte_offset int, mut classes []mir.AbiEightbyteClass, mut visiting map[int]bool) bool {
    mut field_offset := 0
    for field_typ in fields {
        align := m.type_align(field_typ)
        if align > 1 && field_offset % align != 0 {
            field_offset = (field_offset + align - 1) & ~(align - 1)
        }
        if !sysv_classify_type_into(m, field_typ, byte_offset + field_offset, mut classes, mut
            visiting) {
            return false
        }
        field_offset += m.type_size(field_typ)
    }
    return true
}

fn sysv_merge_class_range(mut classes []mir.AbiEightbyteClass, byte_offset int, size int, class mir.AbiEightbyteClass) bool {
    if size <= 0 {
        return true
    }
    start := byte_offset / 8
    end := (byte_offset + size + 7) / 8
    if start < 0 || end > classes.len {
        return false
    }
    for i := start; i < end; i++ {
        classes[i] = sysv_merge_eightbyte_class(classes[i], class)
        if classes[i] == .memory {
            return false
        }
    }
    return true
}

fn sysv_merge_eightbyte_class(a mir.AbiEightbyteClass, b mir.AbiEightbyteClass) mir.AbiEightbyteClass {
    if a == b {
        return a
    }
    if a == .no_class {
        return b
    }
    if b == .no_class {
        return a
    }
    if a == .memory || b == .memory {
        return .memory
    }
    if a == .integer || b == .integer {
        return .integer
    }
    return .sse
}

fn sysv_post_merge_classes(size int, input []mir.AbiEightbyteClass) []mir.AbiEightbyteClass {
    mut classes := input.clone()
    for class in classes {
        if class == .memory {
            return [.memory]
        }
    }
    if size > sysv_max_direct_aggregate_size {
        mut vector_like := classes.len > 0 && classes[0] == .sse
        for i := 1; i < classes.len; i++ {
            if classes[i] != .sseup {
                vector_like = false
                break
            }
        }
        if !vector_like {
            return [.memory]
        }
    }
    for i, class in classes {
        if class == .sseup && (i == 0 || classes[i - 1] !in [.sse, .sseup]) {
            classes[i] = .sse
        }
    }
    return classes
}

fn sysv_assign_value_layout(value_class mir.AbiValueClass, mut state SysVLocationState) mir.AbiValueLayout {
    if value_class.classes.len == 0 {
        return mir.AbiValueLayout{
            value_class: value_class
        }
    }
    if value_class.mode == .indirect || value_class.classes == [mir.AbiEightbyteClass.memory] {
        return sysv_indirect_pointer_layout(value_class, mut state)
    }

    mut needed_int := 0
    mut needed_sse := 0
    for class in value_class.classes {
        match class {
            .integer { needed_int++ }
            .sse { needed_sse++ }
            else {}
        }
    }
    if state.int_regs + needed_int > sysv_int_arg_reg_count
        || state.sse_regs + needed_sse > sysv_sse_arg_reg_count {
        return sysv_stack_value_layout(value_class, mut state)
    }

    mut locs := []mir.AbiLocation{cap: value_class.classes.len}
    mut last_sse := -1
    for i, class in value_class.classes {
        offset := i * 8
        match class {
            .integer {
                locs << mir.AbiLocation{
                    kind:   .int_reg
                    index:  state.int_regs
                    offset: offset
                    class:  class
                }
                state.int_regs++
            }
            .sse {
                last_sse = state.sse_regs
                locs << mir.AbiLocation{
                    kind:   .sse_reg
                    index:  state.sse_regs
                    offset: offset
                    class:  class
                }
                state.sse_regs++
            }
            .sseup {
                locs << mir.AbiLocation{
                    kind:   .sse_reg
                    index:  if last_sse >= 0 { last_sse } else { state.sse_regs }
                    offset: offset
                    class:  class
                }
            }
            else {
                locs << mir.AbiLocation{
                    kind:   .none
                    index:  -1
                    offset: offset
                    class:  class
                }
            }
        }
    }
    return mir.AbiValueLayout{
        value_class: value_class
        locs:        locs
    }
}

fn sysv_indirect_pointer_layout(value_class mir.AbiValueClass, mut state SysVLocationState) mir.AbiValueLayout {
    mut locs := []mir.AbiLocation{cap: 1}
    if state.int_regs < sysv_int_arg_reg_count {
        locs << mir.AbiLocation{
            kind:   .int_reg
            index:  state.int_regs
            offset: 0
            class:  .integer
        }
        state.int_regs++
    } else {
        locs << mir.AbiLocation{
            kind:   .stack
            index:  state.stack_slots
            offset: 0
            class:  .integer
        }
        state.stack_slots++
    }
    return mir.AbiValueLayout{
        value_class: value_class
        locs:        locs
    }
}

fn sysv_stack_value_layout(value_class mir.AbiValueClass, mut state SysVLocationState) mir.AbiValueLayout {
    slots := if value_class.size > 0 { (value_class.size + 7) / 8 } else { value_class.classes.len }
    mut locs := []mir.AbiLocation{cap: slots}
    for i := 0; i < slots; i++ {
        class := if i < value_class.classes.len {
            value_class.classes[i]
        } else {
            mir.AbiEightbyteClass.memory
        }
        locs << mir.AbiLocation{
            kind:   .stack
            index:  state.stack_slots
            offset: i * 8
            class:  class
        }
        state.stack_slots++
    }
    return mir.AbiValueLayout{
        value_class: value_class
        locs:        locs
    }
}

fn fallback_arg_type(m mir.Module, arg_id int) int {
    if arg_id < 0 || arg_id >= m.values.len {
        return 0
    }
    if logical_typ := logical_arg_type_from_value(m, arg_id, 0) {
        return logical_typ
    }
    return m.values[arg_id].typ
}

fn logical_arg_type_from_value(m mir.Module, val_id int, depth int) ?int {
    if depth > 8 || val_id < 0 || val_id >= m.values.len {
        return none
    }
    val := m.values[val_id]
    if val.kind != .instruction {
        return none
    }
    instr := m.instrs[val.index]
    match instr.op {
        .alloca {
            if val.typ <= 0 || val.typ >= m.type_store.types.len {
                return none
            }
            typ := m.type_store.types[val.typ]
            if typ.kind != .ptr_t || typ.elem_type <= 0 || typ.elem_type >= m.type_store.types.len {
                return none
            }
            elem_typ := m.type_store.types[typ.elem_type]
            if elem_typ.kind in [.struct_t, .array_t] {
                return typ.elem_type
            }
        }
        .bitcast {
            if instr.operands.len > 0 {
                return logical_arg_type_from_value(m, instr.operands[0], depth + 1)
            }
        }
        .assign {
            if instr.operands.len > 1 {
                return logical_arg_type_from_value(m, instr.operands[1], depth + 1)
            }
        }
        else {}
    }

    return none
}

fn lower_calls(mut m mir.Module, arch pref.Arch, x64_abi X64Abi, fn_by_name map[string]int) {
    if arch !in [.arm64, .x64] {
        return
    }
    is_x64 := arch == .x64

    for i := 0; i < m.instrs.len; i++ {
        mut instr := &m.instrs[i]
        if instr.op !in [.call, .call_indirect, .call_sret] || instr.operands.len == 0 {
            continue
        }
        ret_typ, sig_param_types := call_signature(m, instr, fn_by_name)
        ret_class := if is_x64 {
            abi_value_class(m, ret_typ, arch, x64_abi)
        } else {
            mir.AbiValueClass{}
        }
        ret_indirect := if is_x64 {
            abi_class_is_indirect(ret_class, m, ret_typ, arch, x64_abi)
        } else {
            needs_indirect(m, ret_typ, arch, x64_abi)
        }
        num_args := instr.operands.len - 1
        instr.abi_arg_class = []mir.AbiArgClass{len: num_args, init: .in_reg}
        if is_x64 {
            instr.abi_arg_classes = []mir.AbiValueClass{len: num_args}
            instr.abi_arg_layouts = []mir.AbiValueLayout{len: num_args}
        }
        mut arg_loc_state := SysVLocationState{}
        if arch == .x64 && x64_abi == .sysv && ret_indirect {
            arg_loc_state.int_regs = 1
        }
        for arg_idx := 0; arg_idx < num_args; arg_idx++ {
            mut arg_typ := 0
            if arg_idx < sig_param_types.len && sig_param_types[arg_idx] > 0 {
                arg_typ = sig_param_types[arg_idx]
            } else {
                arg_id := instr.operands[arg_idx + 1]
                arg_typ = fallback_arg_type(m, arg_id)
            }
            if is_x64 {
                arg_class := abi_value_class(m, arg_typ, arch, x64_abi)
                instr.abi_arg_classes[arg_idx] = arg_class
                if x64_abi == .sysv {
                    instr.abi_arg_layouts[arg_idx] = sysv_assign_value_layout(arg_class, mut
                        arg_loc_state)
                }
                if abi_class_is_indirect(arg_class, m, arg_typ, arch, x64_abi) {
                    instr.abi_arg_class[arg_idx] = .indirect
                }
            } else if needs_indirect(m, arg_typ, arch, x64_abi) {
                instr.abi_arg_class[arg_idx] = .indirect
            }
        }

        instr.abi_ret_class = ret_class
        instr.abi_ret_indirect = ret_indirect
        // Lower ABI-indirect returns to call_sret for backend consumption.
        if instr.abi_ret_indirect {
            instr.op = .call_sret
        }
    }
}

fn call_signature(m mir.Module, instr &mir.Instruction, fn_by_name map[string]int) (int, []int) {
    mut ret_typ := instr.typ
    mut param_types := []int{}
    if instr.operands.len == 0 {
        return ret_typ, param_types
    }

    if instr.op in [.call, .call_sret] {
        callee_id := instr.operands[0]
        if callee_id >= 0 && callee_id < m.values.len {
            callee_name := m.values[callee_id].name
            if callee_name != '' && callee_name in fn_by_name {
                fn_idx := fn_by_name[callee_name]
                if fn_idx >= 0 && fn_idx < m.funcs.len {
                    callee := m.funcs[fn_idx]
                    ret_typ = callee.typ
                    param_types = []int{len: callee.params.len}
                    for i, pid in callee.params {
                        if pid >= 0 && pid < m.values.len {
                            param_types[i] = m.values[pid].typ
                        }
                    }
                }
            }
        }
    } else if instr.op == .call_indirect {
        callee_id := instr.operands[0]
        if callee_id >= 0 && callee_id < m.values.len {
            fn_ptr_typ_id := m.values[callee_id].typ
            if fn_ptr_typ_id > 0 && fn_ptr_typ_id < m.type_store.types.len {
                fn_ptr_typ := m.type_store.types[fn_ptr_typ_id]
                if fn_ptr_typ.kind == .ptr_t && fn_ptr_typ.elem_type > 0
                    && fn_ptr_typ.elem_type < m.type_store.types.len {
                    fn_typ := m.type_store.types[fn_ptr_typ.elem_type]
                    if fn_typ.kind == .func_t {
                        ret_typ = fn_typ.ret_type
                        param_types = fn_typ.params.clone()
                    }
                } else if fn_ptr_typ.kind == .func_t {
                    // Function pointer extracted from struct field via extractvalue
                    // has func_t type directly (not wrapped in ptr_t).
                    ret_typ = fn_ptr_typ.ret_type
                    param_types = fn_ptr_typ.params.clone()
                }
            }
        }
    }

    // Fallback to call operand types when we could not resolve a signature.
    if param_types.len == 0 {
        num_args := if instr.operands.len > 0 { instr.operands.len - 1 } else { 0 }
        param_types = []int{len: num_args}
        for i := 0; i < num_args; i++ {
            arg_id := instr.operands[i + 1]
            param_types[i] = fallback_arg_type(m, arg_id)
        }
    }
    return ret_typ, param_types
}