import "@std/mem"
import "@std/arraylist"
import "@std/debug"
import "@std/enummap"
import "@source/diagnostics"

FunctionId :: distinct u32
GlobalId :: distinct u32
LocalId :: distinct u32
ValueId :: distinct u32
BlockId :: distinct u32
StringId :: distinct u32
ArrayId :: distinct u32
no_array :: maxval!(ArrayId)
no_value :: maxval!(ValueId)

Type :: enum {
    invalid
    int
    float
    bool
    string
    void
    function
    array
    record
    address
    pointer
    i8, i16, i32, isize
    u8, u16, u32, u64, usize
}

aggregate proc(type Type) bool {
    return type == .array or type == .record
}

@hide:file
integer_widths :: $enummap.init_default(Type, usize, usize(0), {
    i8 = 8, u8 = 8,
    i16 = 16, u16 = 16,
    i32 = 32, u32 = 32,
    int = 64, isize = 64, usize = 64, u64 = 64,
})

#! integer storage widths for the supported 64-bit target; zero for noninteger types.
integer_bits proc(type Type) usize {
    if (enummap.get(&integer_widths, type)) |width| return width
    return 0
}

unsigned_integer proc(type Type) bool {
    return type == .u8 or type == .u16 or type == .u32 or type == .u64 or type == .usize
}

#! maxima for supported integer storage types; noninteger keys are absent.
integer_maxima :: $enummap.init_default(Type, u64, u64(0), {
    i8 = 127,
    i16 = 32767,
    i32 = 2147483647,
    int = u64(maxval!(i64)), isize = u64(maxval!(i64)),
    u8 = 255,
    u16 = 65535,
    u32 = 4294967295,
    u64 = maxval!(u64), usize = maxval!(u64),
})

#! minima for supported integer storage types; noninteger keys are absent.
integer_minima :: $enummap.init_default(Type, i64, i64(0), {
    i8 = i64(-128),
    i16 = i64(-32768),
    i32 = i64(-2147483648),
    int = minval!(i64), isize = minval!(i64),
    u8 = 0, u16 = 0, u32 = 0, u64 = 0, usize = 0,
})

ValueData :: union {
    integer i64
    # source pointers use run-wide handles; internal addresses use frame-local IDs below.
    unsigned u64
    real f64
    boolean bool
    string StringId
    function FunctionId
    elements []mut Value
    address u32
}
Value :: struct {
    type Type = .void
    array_type ArrayId = no_array
    raw bool = false
    # interpreter poison survives copies; native storage uses LLVM poison directly.
    uninitialized bool = false
    data ValueData = ValueData{ integer = 0 }
}

RecordField :: struct {
    type Type
    array_type ArrayId = no_array
    offset usize = 0
    byte_offset usize = 0
    tag Value = Value{}
}

#! fixed recursive aggregate shape; cells count interpreter slots, byte_size describes native layout.
#! unions retain overlapping byte cells; ordinary aggregates retain flattened scalar leaves.
ArrayType :: struct {
    kind Type = .array
    element Type = .invalid
    element_array ArrayId = no_array
    len usize = 0
    cells usize = 0
    mutable bool = false
    fields []mut RecordField = &[]
    is_union bool = false
    tag_type Type = .void
    payload_offset usize = 0
    byte_size usize = 0
    alignment usize = 1
}

#! maps a local slot or temporary to its preallocated, by-value frame segment.
ArrayStorage :: struct {
    index u32
    array_type ArrayId
    offset usize
    local bool
}

Op :: enum {
    constant
    undefined_value
    load_local
    store_local
    assign_local
    reset_local
    load_global
    array_literal
    record_literal
    union_literal
    union_tag
    field
    index
    address_local
    address_global
    index_address
    field_address
    load_address
    store_address
    pointer_from_address
    pointer_address
    pointer_offset
    pointer_cast
    string_pointer
    string_length
    slice_check
    sentinel_check
    optional_check
    memory_copy, memory_set
    negate
    convert
    logical_not
    bit_not, bit_and, bit_or, bit_xor
    shift_left, shift_right, shift_left_saturating
    add, subtract
    multiply, divide
    call
    equal, not_equal
    less, less_equal
    greater, greater_equal
}

Instruction :: struct {
    op Op
    type Type = .void
    # pointer operations describe their pointee here, not aggregate result storage.
    array_type ArrayId = no_array
    pointee Type = .invalid
    span diagnostics.Span = diagnostics.Span{ start = 0, end = 0 }
    destination ValueId = no_value
    left ValueId = no_value
    right ValueId = no_value
    extra ValueId = no_value
    slot u32 = 0
    constant Value = Value{}
    arguments_start usize = 0
    arguments_count usize = 0
}

TerminatorKind :: enum {
    return_value,
    trap,
    jump,
    branch,
    unreachable,
}

Terminator :: struct {
    kind TerminatorKind = .return_value
    value ValueId = no_value
    target BlockId = BlockId(0)
    alternative BlockId = BlockId(0)
    diagnostic diagnostics.DiagnosticId = diagnostics.no_diagnostic
    span diagnostics.Span = diagnostics.Span{ start = 0, end = 0 }
}

Block :: struct {
    instructions arraylist.ArrayList(Instruction)
    arguments arraylist.ArrayList(ValueId)
    terminator Terminator = Terminator{}
}

#! construction derives storage segments, aggregate totals, value counts, and address presence.
#! producers use add_local, add_parameter, and emit; derived metadata is read-only to callers.
Function :: struct {
    parameters arraylist.ArrayList(Type)
    # owned by the module allocator; empty for interpreted/native-defined functions.
    external_name []u8 = ""
    locals arraylist.ArrayList(Type)
    result Type = .void
    result_array_type ArrayId = no_array
    result_fallible bool = false
    array_storage arraylist.ArrayList(ArrayStorage)
    aggregate_cells usize = 0
    has_addresses bool = false
    blocks arraylist.ArrayList(Block)
    entry BlockId = BlockId(0)
    value_count usize = 0
    span diagnostics.Span = diagnostics.Span{ start = 0, end = 0 }
}

Global :: struct {
    type Type
    array_type ArrayId = no_array
    initializer FunctionId
    span diagnostics.Span
}

#! retains the original allocation for freeing while exposing only decoded bytes.
StringStorage :: struct {
    allocation []mut u8
    len usize
}

#! owns executable storage, including decoded strings; no AST or binding table is retained.
#! spans and diagnostic IDs refer to the caller's source and diagnostic store.
Module :: struct {
    allocator mem.Allocator
    functions arraylist.ArrayList(Function)
    globals arraylist.ArrayList(Global)
    strings arraylist.ArrayList(StringStorage)
    arrays arraylist.ArrayList(ArrayType)
    entry ?FunctionId = null
    entry_diagnostic diagnostics.DiagnosticId = diagnostics.no_diagnostic
}

#! retains `allocator` for executable storage; release the returned module with deinit.
init proc(allocator mem.Allocator) Module {
    return Module{
        allocator = allocator,
        functions = arraylist.init(allocator),
        globals = arraylist.init(allocator),
        strings = arraylist.init(allocator),
        arrays = arraylist.init(allocator),
    }
}

#! retains `allocator` in its lists; transfer the function to a module using the same allocator for cleanup with deinit.
init_function proc(allocator mem.Allocator) Function {
    return Function{
        parameters = arraylist.init(allocator),
        locals = arraylist.init(allocator),
        blocks = arraylist.init(allocator),
        array_storage = arraylist.init(allocator),
    }
}

#! retains `allocator` in its lists; transfer the block to a module function for cleanup with deinit.
init_block proc(allocator mem.Allocator) Block {
    return Block{
        instructions = arraylist.init(allocator),
        arguments = arraylist.init(allocator),
    }
}

#! adds a fixed shape and derives its flattened storage size; child shapes must already exist.
#! scalar elements use no_array; zero-leaf children retain their logical lengths.
add_array proc(
    module @mut Module,
    element Type,
    element_array ArrayId,
    len usize,
    mutable bool,
) ArrayId ! (mem.AllocError | diagnostics.CapacityError) {
    debug.assert(aggregate(element) == (element_array != no_array))
    width :: if (aggregate(element)) module.arrays.items[usize(element_array)].cells else 1
    if (width != 0 and len > divtrunc!(maxval!(usize), width)) return .storage_limit_reached
    cells :: len * width
    if (cells > divtrunc!(maxval!(usize), sizeof!(Value))) return .storage_limit_reached
    if (module.arrays.items.len >= usize(no_array)) return .storage_limit_reached
    stride :: storage_size(module, element, element_array)
    if (stride != 0 and len > divtrunc!(maxval!(usize), stride)) return .storage_limit_reached
    id :: ArrayId(module.arrays.items.len)
    try arraylist.append(&module.arrays, ArrayType{
        element = element,
        element_array = element_array,
        len = len,
        cells = cells,
        mutable = mutable,
        byte_size = len * stride,
        alignment = storage_alignment(module, element, element_array),
    })
    return id
}

#! copies declaration-order fields and derives their flattened offsets; child shapes must exist.
add_record proc(
    module @mut Module,
    fields []RecordField,
) ArrayId ! (mem.AllocError | diagnostics.CapacityError) {
    if (module.arrays.items.len >= usize(no_array)) return .storage_limit_reached
    owned :: try mem.alloc(RecordField, module.allocator, fields.len)
    errdefer mem.free(module.allocator, owned)
    cells usize := 0
    bytes usize := 0
    alignment usize := 1
    for (fields) |field, index| {
        debug.assert(aggregate(field.type) == (field.array_type != no_array))
        width :: if (aggregate(field.type)) module.arrays.items[usize(field.array_type)].cells else 1
        if (width > divtrunc!(maxval!(usize), sizeof!(Value)) - cells) return .storage_limit_reached
        field_alignment :: storage_alignment(module, field.type, field.array_type)
        if (field_alignment > alignment) alignment = field_alignment
        if (bytes > maxval!(usize) - (field_alignment - 1)) return .storage_limit_reached
        bytes = aligned_size(bytes, field_alignment)
        size :: storage_size(module, field.type, field.array_type)
        if (size > maxval!(usize) - bytes) return .storage_limit_reached
        owned[index] = RecordField{
            type = field.type, array_type = field.array_type, offset = cells,
            byte_offset = bytes, tag = field.tag,
        }
        bytes += size
        cells += width
    }
    id :: ArrayId(module.arrays.items.len)
    if (bytes > maxval!(usize) - (alignment - 1)) return .storage_limit_reached
    try arraylist.append(&module.arrays, ArrayType{
        kind = .record, fields = owned, cells = cells,
        byte_size = aligned_size(bytes, alignment), alignment = alignment,
    })
    return id
}

#! returns a borrowed child, including a raw view when projecting overlapping union storage.
aggregate_item proc(module @Module, value Value, index usize) Value {
    shape :: module.arrays.items[usize(value.array_type)]
    type := shape.element
    child := shape.element_array
    start usize := 0
    width usize := 1
    if (shape.kind == .record) {
        field :: shape.fields[index]
        type = field.type
        child = field.array_type
        start = field.offset
        if (value.raw) start = field.byte_offset
        if (shape.is_union) start = shape.payload_offset
    } else {
        if (value.raw) width = storage_size(module, type, child)
        else if (aggregate(type)) width = module.arrays.items[usize(child)].cells
        start = index * width
    }
    raw :: value.raw or shape.is_union
    if (!aggregate(type)) {
        if (raw) return read_storage(module, type, child, value.data.elements[start..], true)
        if (value.data.elements[start].uninitialized) return Value{ type = type, uninitialized = true }
        return value.data.elements[start]
    }
    width = if (raw) storage_size(module, type, child) else module.arrays.items[usize(child)].cells
    return Value{
        type = type, array_type = child, raw = raw,
        data = ValueData{ elements = value.data.elements[start..start + width] },
    }
}

#! copies variant descriptors; all payloads share one eight-byte-aligned storage region.
add_union proc(module @mut Module, fields []RecordField, tagged bool) ArrayId ! (mem.AllocError | diagnostics.CapacityError) {
    if (module.arrays.items.len >= usize(no_array)) return .storage_limit_reached
    owned :: try mem.alloc(RecordField, module.allocator, fields.len)
    errdefer mem.free(module.allocator, owned)
    size usize := 0
    for (fields) |field, index| {
        width :: storage_size(module, field.type, field.array_type)
        if (width > size) size = width
        owned[index] = field
        owned[index].offset = if (tagged) 8 else 0
        owned[index].byte_offset = owned[index].offset
    }
    prefix usize :: if (tagged) 8 else 0
    if (size > divtrunc!(maxval!(usize), sizeof!(Value)) - 15) return .storage_limit_reached
    size = aligned_size(size, 8) + prefix
    id :: ArrayId(module.arrays.items.len)
    tag_type :: if (tagged and fields.len != 0) fields[0].tag.type else Type.void
    try arraylist.append(&module.arrays, ArrayType{
        kind = .record, fields = owned, is_union = true,
        tag_type = tag_type,
        payload_offset = prefix, cells = size, byte_size = size, alignment = 8,
    })
    return id
}

@hide:file
aligned_size proc(size, alignment usize) usize {
    return divtrunc!(size + alignment - 1, alignment) * alignment
}

@hide:file
scalar_alignments :: $enummap.init_default(Type, usize, usize(1), {
    float = 8, string = 8, function = 8, pointer = 8,
})

@hide:file
scalar_sizes :: $enummap.init_default(Type, usize, usize(0), {
    float = 8, function = 8, pointer = 8,
    string = 16,
    bool = 1,
    void = 0,
})

storage_alignment proc(module @Module, type Type, shape ArrayId) usize {
    if (aggregate(type)) return module.arrays.items[usize(shape)].alignment
    bits :: integer_bits(type)
    if (bits != 0) return divtrunc!(bits, 8)
    if (enummap.get(&scalar_alignments, type)) |alignment| return alignment
    return 1
}

storage_size proc(module @Module, type Type, shape ArrayId) usize {
    if (aggregate(type)) return module.arrays.items[usize(shape)].byte_size
    bits :: integer_bits(type)
    if (bits != 0) return divtrunc!(bits, 8)
    return enummap.get(&scalar_sizes, type)?
}

#! loads a scalar or borrows an aggregate from byte cells without allocating.
read_storage proc(module @Module, type Type, shape ArrayId, cells []mut Value, raw bool) Value {
    if (aggregate(type)) return Value{ type = type, array_type = shape, raw = raw, data = ValueData{ elements = cells[..storage_size(module, type, shape)] } }
    bits u64 := 0
    size :: storage_size(module, type, shape)
    count :: if (size > 8) usize(8) else size
    for (0..count) |i| {
        if (cells[i].uninitialized) return Value{ type = type, uninitialized = true }
    }
    for (0..count) |i| bits += cells[i].data.unsigned << u64(i * 8)
    data := ValueData{ unsigned = bits }
    match type {
        .i8: data = ValueData{ integer = i64(ptrcast!(@i8, &bits)^) }
        .i16: data = ValueData{ integer = i64(ptrcast!(@i16, &bits)^) }
        .i32: data = ValueData{ integer = i64(ptrcast!(@i32, &bits)^) }
        .int, .isize: data = ValueData{ integer = ptrcast!(@i64, &bits)^ }
        .float: data = ValueData{ real = ptrcast!(@f64, &bits)^ }
        .bool: data = ValueData{ boolean = (bits & 1) != 0 }
        .string: data = ValueData{ string = StringId(bits) }
        .function: data = ValueData{ unsigned = bits }
        .void: return Value{}
        else: {}
    }
    return Value{ type = type, data = data }
}

#! stores a value in shared byte cells; recursive aggregates preserve their native field padding.
write_storage proc(module @Module, cells []mut Value, value Value) void {
    if (aggregate(value.type)) {
        shape :: module.arrays.items[usize(value.array_type)]
        if (value.raw or shape.is_union) {
            memcopy!(cells[..shape.byte_size], value.data.elements)
            return
        }
        count :: if (shape.kind == .record) shape.fields.len else shape.len
        for (0..count) |i| {
            offset :: if (shape.kind == .record) shape.fields[i].byte_offset else i * storage_size(module, shape.element, shape.element_array)
            write_storage(module, cells[offset..], aggregate_item(module, value, i))
        }
        return
    }
    bits u64 := 0
    match value.type {
        .float: bits = ptrcast!(@u64, &value.data.real)^
        .bool: bits = if (value.data.boolean) 1 else 0
        .string: bits = u64(value.data.string)
        .function: bits = if (value.data.function == maxval!(FunctionId)) maxval!(u64) else u64(value.data.function)
        .pointer: bits = value.data.unsigned
        .void: return
        else: bits = ptrcast!(@u64, &value.data.integer)^
    }
    size :: storage_size(module, value.type, value.array_type)
    for (0..size) |i| {
        byte u64 := 0
        if (i < 8) byte = (bits >> u64(i * 8)) & 255
        else byte = (u64(string_bytes(module, value.data.string).len) >> u64((i - 8) * 8)) & 255
        cells[i] = Value{ type = .u8, uninitialized = value.uninitialized, data = ValueData{ unsigned = byte } }
    }
}

#! converts raw aggregate projections into the caller's preallocated flattened segment.
copy_aggregate proc(module @Module, destination Value, source Value) void {
    shape :: module.arrays.items[usize(source.array_type)]
    if (!source.raw or shape.is_union) {
        memcopy!(destination.data.elements, source.data.elements)
        return
    }
    count :: if (shape.kind == .record) shape.fields.len else shape.len
    for (0..count) |i| {
        child :: aggregate_item(module, source, i)
        if (aggregate(child.type)) copy_aggregate(module, aggregate_item(module, destination, i), child)
        else {
            offset :: if (shape.kind == .record) shape.fields[i].offset else i
            destination.data.elements[offset] = child
        }
    }
}

union_tag proc(module @Module, value Value) Value {
    shape :: module.arrays.items[usize(value.array_type)]
    return read_storage(module, shape.tag_type, no_array, value.data.elements, true)
}

union_active proc(module @Module, value Value, field usize) bool {
    shape :: module.arrays.items[usize(value.array_type)]
    if (shape.tag_type == .void) return true
    tag :: union_tag(module, value)
    return tag.data.unsigned == shape.fields[field].tag.data.unsigned
}

@hide:file
prepare_storage proc(module @Module, function @mut Function, index u32, shape ArrayId, local bool) void ! (mem.AllocError | diagnostics.CapacityError) {
    cells :: module.arrays.items[usize(shape)].cells
    if (cells > divtrunc!(maxval!(usize), sizeof!(Value)) - function.aggregate_cells) return .storage_limit_reached
    try arraylist.append(&function.array_storage, ArrayStorage{
        index = index,
        array_type = shape,
        offset = function.aggregate_cells,
        local = local,
    })
    function.aggregate_cells += cells
}

#! creates a local and its disjoint by-value backing together; scalar locals use no_array.
add_local proc(
    module @mut Module,
    function FunctionId,
    type Type,
    shape ArrayId,
) LocalId ! (mem.AllocError | diagnostics.CapacityError) {
    debug.assert(aggregate(type) == (shape != no_array))
    body :: &module.functions.items[usize(function)]
    if (body.locals.items.len >= usize(maxval!(LocalId))) return .storage_limit_reached
    id :: LocalId(body.locals.items.len)
    try arraylist.append(&body.locals, type)
    errdefer _ = arraylist.pop(&body.locals)
    if (aggregate(type)) try prepare_storage(module, body, u32(id), shape, true)
    return id
}

#! creates a parameter and its local backing; all parameters precede non-parameter locals.
add_parameter proc(
    module @mut Module,
    function FunctionId,
    type Type,
    shape ArrayId,
) LocalId ! (mem.AllocError | diagnostics.CapacityError) {
    body :: &module.functions.items[usize(function)]
    debug.assert(body.parameters.items.len == body.locals.items.len)
    try arraylist.append(&body.parameters, type)
    errdefer _ = arraylist.pop(&body.parameters)
    return try add_local(module, function, type, shape)
}

#! appends an instruction and derives its result count, aggregate backing, and address presence.
#! no_value requests a fresh result; explicit result IDs support non-SSA definitions across blocks.
#! repeated aggregate results must retain their shape and reuse their existing segment.
#! failures leave instruction and storage lengths, totals, counts, and flags unchanged.
emit proc(
    module @mut Module,
    function FunctionId,
    block BlockId,
    instruction Instruction,
) ValueId ! (mem.AllocError | diagnostics.CapacityError) {
    body :: &module.functions.items[usize(function)]
    target :: &body.blocks.items[usize(block)]
    operation := instruction
    produces_value :: (
        operation.op != .store_local
        and operation.op != .assign_local
        and operation.op != .reset_local
        and operation.op != .store_address
        and operation.op != .slice_check
        and operation.op != .sentinel_check
        and operation.op != .optional_check
        and operation.op != .memory_copy
        and operation.op != .memory_set
    )
    count := body.value_count
    if (produces_value) {
        if (operation.destination == no_value) {
            if (count >= usize(no_value)) return .storage_limit_reached
            operation.destination = ValueId(count)
        }
        if (usize(operation.destination) >= count) count = usize(operation.destination) + 1
    } else {
        debug.assert(operation.destination == no_value)
    }
    try arraylist.append(&target.instructions, operation)
    errdefer _ = arraylist.pop(&target.instructions)
    if (produces_value and aggregate(operation.type)) {
        debug.assert(operation.array_type != no_array)
        registered := false
        if (usize(operation.destination) < body.value_count) {
            for (body.array_storage.items) |storage| {
                if (storage.local or storage.index != u32(operation.destination)) continue
                debug.assert(storage.array_type == operation.array_type)
                registered = true
                break
            }
        }
        if (!registered) try prepare_storage(module, body, u32(operation.destination), operation.array_type, false)
    }
    body.value_count = count
    if (operation.type == .address) body.has_addresses = true
    return operation.destination
}

#! resolves a module-local ID to bytes borrowed until module deinit.
string_bytes proc(module @Module, id StringId) []u8 {
    string :: module.strings.items[usize(id)]
    return string.allocation[..string.len]
}

#! releases all executable storage using the retained allocator.
deinit proc(module @mut Module) void {
    for (0..(module.functions.items.len)) |i| {
        function :: &module.functions.items[i]
        for (0..(function.blocks.items.len)) |j| {
            block :: &function.blocks.items[j]
            arraylist.deinit(&block.instructions)
            arraylist.deinit(&block.arguments)
        }
        arraylist.deinit(&function.parameters)
        arraylist.deinit(&function.locals)
        arraylist.deinit(&function.blocks)
        arraylist.deinit(&function.array_storage)
        if (function.external_name.len != 0) mem.free(module.allocator, function.external_name)
    }
    for (module.strings.items) |string|
        mem.free(module.allocator, string.allocation)

    for (module.arrays.items) |shape| {
        if (shape.kind == .record) mem.free(module.allocator, shape.fields)
    }
    arraylist.deinit(&module.functions)
    arraylist.deinit(&module.globals)
    arraylist.deinit(&module.strings)
    arraylist.deinit(&module.arrays)
}