import "@std/mem"
import "@std/arraylist"
import "@std/debug"
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,
    usize,
    address,
}

ValueData :: union {
    integer i64
    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
    data ValueData = ValueData{ integer = 0 }
}

#! fixed recursive shape; add_array derives `cells`, the flattened scalar-leaf count.
ArrayType :: struct {
    element Type
    element_array ArrayId = no_array
    len usize
    cells usize
    mutable bool
}

#! 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
    load_local
    store_local
    assign_local
    reset_local
    load_global
    array_literal
    index
    address_local
    address_global
    index_address
    load_address
    store_address
    negate
    logical_not
    add
    subtract
    multiply
    divide
    call
    equal
    not_equal
    less
    less_equal
    greater
    greater_equal
}

Instruction :: struct {
    op Op
    type Type = .void
    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
    slot u32 = 0
    constant Value = Value{}
    arguments_start usize = 0
    arguments_count usize = 0
}

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

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)
    locals arraylist.ArrayList(Type)
    result Type = .void
    result_array_type ArrayId = no_array
    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
}

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),
    }
}

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),
    }
}

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((element == .array) == (element_array != no_array))
    width :: if (element == .array) 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
    id :: ArrayId(module.arrays.items.len)
    try arraylist.append(&module.arrays, ArrayType{
        element = element,
        element_array = element_array,
        len = len,
        cells = cells,
        mutable = mutable,
    })
    return id
}

@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((type == .array) == (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 (type == .array) 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, array backing, and address presence.
#! no_value requests a fresh result; explicit result IDs support non-SSA definitions across blocks.
#! repeated array 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
    )
    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 operation.type == .array) {
        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)
    }
    for (module.strings.items) |string|
        mem.free(module.allocator, string.allocation)

    arraylist.deinit(&module.functions)
    arraylist.deinit(&module.globals)
    arraylist.deinit(&module.strings)
    arraylist.deinit(&module.arrays)
}