// 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 builder

import v2.ast
import v2.transformer
import runtime

$if !windows {
    fn C.pthread_create(thread &C.pthread_t, attr voidptr, start_routine fn (voidptr) voidptr, arg voidptr) int
    fn C.pthread_join(thread C.pthread_t, retval voidptr) int
    fn C.pthread_attr_init(attr voidptr) int
    fn C.pthread_attr_setstacksize(attr voidptr, stacksize usize) int
    fn C.pthread_attr_destroy(attr voidptr) int

    // TransformChunkFlatArgs carries one worker's slice of the parallel flat
    // transform: the worker emits its file range cursor-native into its own
    // FlatBuilder (no legacy ast.File rehydrate) and hands back the resulting
    // FlatAst, which the main thread merges via `FlatBuilder.append_flat`.
    struct TransformChunkFlatArgs {
        t                voidptr // &transformer.Transformer
        flat             &ast.FlatAst = unsafe { nil }
        flat_extra_stmts [][]ast.Stmt
        flat_start       int
        flat_end         int
        result_ptr       voidptr // &ast.FlatAst
        worker_ptr       voidptr
        worker_idx       int
    }

    fn transform_chunk_flat_thread(arg voidptr) voidptr {
        a := unsafe { &TransformChunkFlatArgs(arg) }
        t := unsafe { &transformer.Transformer(a.t) }
        mut w := t.new_worker_clone(a.worker_idx)
        // Cursor-native: transform one file at a time straight into this
        // worker's FlatBuilder. Peak per worker is the cumulative flat for its
        // chunk (no legacy ast.File ever materialises).
        mut wb := ast.new_flat_builder()
        for fi := a.flat_start; fi < a.flat_end; fi++ {
            extra := if fi >= 0 && fi < a.flat_extra_stmts.len {
                a.flat_extra_stmts[fi]
            } else {
                []ast.Stmt{}
            }
            w.transform_file_index_with_extra_to_flat(a.flat, fi, extra, mut wb)
        }
        unsafe {
            *(&ast.FlatAst(a.result_ptr)) = wb.flat
            *(&voidptr(a.worker_ptr)) = voidptr(w)
        }
        return unsafe { nil }
    }
}

fn flat_extra_stmts_by_file(extra_stmts map[int][]ast.Stmt, n_files int) [][]ast.Stmt {
    mut out := [][]ast.Stmt{cap: n_files}
    for _ in 0 .. n_files {
        out << []ast.Stmt{}
    }
    for fi, stmts in extra_stmts {
        if fi >= 0 && fi < n_files {
            out[fi] = stmts
        }
    }
    return out
}

// transform_files_parallel_flat_direct is the flat-native parallel transform
// used by every backend. It mirrors the sequential
// `transform_flat_to_flat_direct` (pre_pass -> prepare -> per-file cursor
// transform -> post_pass tail) but fans the per-file loop across worker threads:
// each worker emits its contiguous file range cursor-native into its OWN
// FlatBuilder, then the main thread concatenates them in file order via
// `FlatBuilder.append_flat`. No legacy ast.File is materialised; backends that
// still consume []ast.File (.v/eval) rehydrate from the transformed flat at
// the codegen boundary in builder.v.
//
// Per-worker synth-position disjointness is handled by `new_worker_clone`
// (it offsets synth_pos_counter by `-worker_idx * 100_000`); worker state is
// merged back via `merge_worker` after the join, exactly as the legacy path.
fn (mut b Builder) transform_files_parallel_flat_direct(mut trans transformer.Transformer) ast.FlatAst {
    // Pre-pass + generic preparation: sequential, once (same as the legacy path).
    trans.pre_pass_from_flat(&b.flat)
    mut flat_extra_stmts := [][]ast.Stmt{}
    if trans.needs_full_files_for_transform() {
        extra_stmts := trans.prepare_flat_for_transform(&b.flat)
        flat_extra_stmts = flat_extra_stmts_by_file(extra_stmts, b.flat.files.len)
    }

    n_jobs := runtime.nr_jobs()
    n_files := b.flat.files.len
    mut out := ast.new_flat_builder()
    $if windows {
        for fi in 0 .. n_files {
            extra := if fi < flat_extra_stmts.len { flat_extra_stmts[fi] } else { []ast.Stmt{} }
            trans.transform_file_index_with_extra_to_flat(&b.flat, fi, extra, mut out)
        }
    } $else {
        if n_files <= 1 || n_jobs <= 1 {
            for fi in 0 .. n_files {
                extra := if fi < flat_extra_stmts.len { flat_extra_stmts[fi] } else { []ast.Stmt{} }
                trans.transform_file_index_with_extra_to_flat(&b.flat, fi, extra, mut out)
            }
        } else {
            // Contiguous [start,end) file ranges, one per worker. Merging the
            // worker outputs in spawn order preserves original file order.
            mut bucket_indices := [][]int{len: n_jobs}
            chunk_size := (n_files + n_jobs - 1) / n_jobs
            mut i := 0
            mut bw := 0
            for i < n_files {
                end := if i + chunk_size < n_files { i + chunk_size } else { n_files }
                for j in i .. end {
                    bucket_indices[bw] << j
                }
                i = end
                bw++
            }

            mut chunk_flats := []ast.FlatAst{len: n_jobs}
            mut worker_ptrs := []voidptr{len: n_jobs, init: unsafe { nil }}
            mut thread_ids := []C.pthread_t{len: n_jobs}
            mut args := []TransformChunkFlatArgs{cap: n_jobs}

            // 64 MB worker stacks: ARM64-compiled transform recursion is stack-heavy.
            attr_buf := [64]u8{}
            attr := unsafe { voidptr(&attr_buf[0]) }
            C.pthread_attr_init(attr)
            C.pthread_attr_setstacksize(attr, 64 * 1024 * 1024)

            mut chunk_idx := 0
            for w in 0 .. n_jobs {
                idxs := bucket_indices[w]
                if idxs.len == 0 {
                    continue
                }
                args << TransformChunkFlatArgs{
                    t:                unsafe { voidptr(trans) }
                    flat:             unsafe { &b.flat }
                    flat_extra_stmts: flat_extra_stmts
                    flat_start:       idxs[0]
                    flat_end:         idxs[idxs.len - 1] + 1
                    result_ptr:       unsafe { voidptr(&chunk_flats[chunk_idx]) }
                    worker_ptr:       unsafe { voidptr(&worker_ptrs[chunk_idx]) }
                    worker_idx:       chunk_idx
                }
                C.pthread_create(unsafe { &thread_ids[chunk_idx] }, attr,
                    transform_chunk_flat_thread, unsafe { voidptr(&args[chunk_idx]) })
                chunk_idx++
            }
            C.pthread_attr_destroy(attr)
            for ci := 0; ci < chunk_idx; ci++ {
                C.pthread_join(thread_ids[ci], unsafe { nil })
            }

            // Concatenate worker flats in spawn (= file) order, merging worker state.
            mut ci := 0
            for w in 0 .. n_jobs {
                if bucket_indices[w].len == 0 {
                    continue
                }
                out.append_flat(chunk_flats[ci])
                worker := unsafe { &transformer.Transformer(worker_ptrs[ci]) }
                trans.merge_worker(worker)
                ci++
            }
            // Move synth_pos_counter past all worker ranges for the post_pass.
            trans.set_synth_pos_counter(-(chunk_idx * 100_000) - 1)
        }
    }

    // Post-pass tail: identical to transform_flat_to_flat_direct's tail.
    generated_parts := trans.generated_fns_parts_from_flat(&out.flat)
    trans.post_pass_to_flat(mut out, generated_parts)
    trans.apply_post_pass_tail_from_flat(&out.flat)
    return out.flat
}