// Copyright (c) 2019-2024 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
//
// GMP Scheduler - translated from Go's runtime/proc.go.
//
// The scheduler distributes ready-to-run goroutines over worker threads.
// Key functions translated from Go:
//   - schedule()      -> scheduler_loop()
//   - findRunnable()  -> find_runnable()
//   - execute()       -> execute()
//   - newproc()       -> goroutine_create()
//   - wakep()         -> wake_p()
//   - runqput()       -> runq_put()
//   - runqget()       -> runq_get()
//   - stealWork()     -> steal_work()
module goroutines

// import sync as _
// import runtime as _

// goroutine_create creates a new goroutine to run `f` with argument `arg`.
// This is the equivalent of Go's `newproc` function.
// Called by the compiler for `go expr()`.
pub fn goroutine_create(f voidptr, arg voidptr, arg_size int) {
	gp := newproc1(f, arg, arg_size)
	// Put on global run queue so any M can pick it up.
	// This is simpler than local queue + work stealing and avoids
	// visibility issues when the creator is the main (non-scheduler) thread.
	glob_runq_put(gp)
	// Try to wake an idle P if needed
	wake_p()
}

// newproc1 allocates and initializes a new G.
// Translated from Go's newproc1 in proc.go.
fn newproc1(f voidptr, arg voidptr, arg_size int) &Goroutine {
	// Try to get a dead G from the local P's free list first
	mut pp := get_current_p()
	mut gp := if pp != unsafe { nil } {
		gfget(mut pp)
	} else {
		unsafe { &Goroutine(nil) }
	}
	if gp == unsafe { nil } {
		// Try global free list
		gp = gfget_global()
	}
	if gp == unsafe { nil } {
		// Allocate a new G
		gp = &Goroutine{}
	}

	// Allocate or reuse stack
	if gp.stack == unsafe { nil } {
		stack_size := default_stack_size
		gp.stack = unsafe { malloc(stack_size) }
		gp.stack_size = stack_size
	}

	gp.fn_ptr = f
	gp.fn_arg = arg
	gp.status = .runnable
	gp.preempt = false

	// Assign goroutine ID
	gp.id = assign_goid()

	// Initialize context to run goroutine_entry
	context_init(mut &gp.context, gp.stack, gp.stack_size, goroutine_trampoline, voidptr(gp))

	// Track all goroutines
	allgs_mu.acquire()
	allgs << gp
	allgs_mu.release()

	return gp
}

// goroutine_trampoline is the entry point for all goroutines.
// It calls the user function, then cleans up and returns to the scheduler.
// Equivalent to Go's goexit -> goexit0.
fn goroutine_trampoline(arg voidptr) {
	mut gp := unsafe { &Goroutine(arg) }

	// Call the actual function
	if gp.fn_ptr != unsafe { nil } {
		call_goroutine_fn(gp.fn_ptr, gp.fn_arg)
	}

	// Function returned - goroutine is done
	goexit0(mut gp)
}

// call_goroutine_fn calls the goroutine's function pointer with the given argument.
fn call_goroutine_fn(fn_ptr voidptr, arg voidptr) {
	unsafe {
		cb := GoFn(fn_ptr)
		cb(arg)
	}
}

// goexit0 handles goroutine cleanup after the user function returns.
// Translated from Go's goexit0 in proc.go.
fn goexit0(mut gp Goroutine) {
	mut mp := get_current_m()

	gp.status = .dead
	gp.m = unsafe { nil }
	gp.fn_ptr = unsafe { nil }
	gp.fn_arg = unsafe { nil }
	gp.wait_reason = ''
	gp.preempt = false

	// Dissociate G from M
	if mp != unsafe { nil } {
		mp.curg = unsafe { nil }
	}

	// Put the dead G on the free list for reuse
	mut pp := get_current_p()
	if pp != unsafe { nil } {
		gfput(mut pp, gp)
	} else {
		gfput_global(gp)
	}

	// Switch back to the scheduler (g0 context).
	// This returns to schedule_loop via execute's context_switch.
	if mp != unsafe { nil } && mp.g0 != unsafe { nil } {
		context_set(&mp.g0.context)
	}
}

// schedule is the main scheduler entry point.
// Finds a runnable goroutine and executes it. Never returns.
// Translated from Go's schedule() in proc.go.
pub fn schedule() {
	mut mp := get_current_m()
	if mp == unsafe { nil } {
		return
	}
	mut pp := mp.p
	if pp == unsafe { nil } {
		return
	}
	pp.sched_tick++

	// Find a runnable goroutine
	gp, _ := find_runnable(mut mp, mut pp)
	if gp == unsafe { nil } {
		// No work found - park this M
		park_m(mut mp)
		return
	}

	// Execute the goroutine
	execute(mut mp, mut pp, gp)
}

// find_runnable finds a runnable goroutine to execute.
// Tries: local queue, global queue, work stealing from other P's.
// Translated from Go's findRunnable() in proc.go.
fn find_runnable(mut mp Machine, mut pp Processor) (&Goroutine, bool) {
	// Check global queue every Nth tick for fairness (Go uses 61)
	if pp.sched_tick % global_queue_check_interval == 0 {
		gp := glob_runq_get()
		if gp != unsafe { nil } {
			return gp, false
		}
	}

	// Try local run queue first
	gp, inherit := runq_get(mut pp)
	if gp != unsafe { nil } {
		return gp, inherit
	}

	// Try global run queue
	gp2 := glob_runq_get()
	if gp2 != unsafe { nil } {
		return gp2, false
	}

	// Try to steal work from other P's
	gp3 := steal_work(mut pp)
	if gp3 != unsafe { nil } {
		return gp3, false
	}

	return unsafe { nil }, false
}

// execute runs a goroutine on the current M.
// Translated from Go's execute() in proc.go.
fn execute(mut mp Machine, mut pp Processor, gp &Goroutine) {
	mut g := unsafe { gp }
	// Associate G with M
	mp.curg = g
	g.m = mp
	g.status = .running

	// Switch context to the goroutine
	context_switch(mut &mp.g0.context, &g.context)
	// When we return here, the goroutine has yielded back to us
	// The scheduler loop will be re-entered
}

// wake_p tries to wake an idle P to run goroutines.
// Translated from Go's wakep() in proc.go.
fn wake_p() {
	// Check if there's an idle P
	gsched.mu.acquire()
	if gsched.npidle == 0 {
		gsched.mu.release()
		return
	}
	// Don't wake if there are already spinning M's
	if gsched.nmspinning > 0 {
		gsched.mu.release()
		return
	}
	// Get an idle P
	pp := pid_get()
	if pp == unsafe { nil } {
		gsched.mu.release()
		return
	}
	gsched.mu.release()

	// Start a new M for this P (or wake an idle one)
	start_m(pp)
}

// park_m parks the current M - it goes to sleep waiting for work.
fn park_m(mut mp Machine) {
	// Release P and add M to idle list under lock
	gsched.mu.acquire()
	if mp.p != unsafe { nil } {
		pid_put(mp.p)
		mp.p = unsafe { nil }
	}
	mp.sched_link = gsched.midle
	gsched.midle = mp
	gsched.nmidle++
	gsched.mu.release()

	// Sleep until woken
	mp.park.wait()

	// Woken up - acquire a P and return to schedule_loop
	acquire_p(mut mp)
}

// acquire_p tries to get an idle P for the given M.
fn acquire_p(mut mp Machine) {
	gsched.mu.acquire()
	mut pp := pid_get()
	gsched.mu.release()
	if pp != unsafe { nil } {
		wire_p(mut mp, mut pp)
	}
}

// wire_p associates a P with an M.
fn wire_p(mut mp Machine, mut pp Processor) {
	mp.p = pp
	pp.m = mp
	pp.status = .running
}

// start_m starts or wakes an M to run the given P.
// Translated from Go's startm() in proc.go.
fn start_m(pp &Processor) {
	gsched.mu.acquire()

	// Try to get an idle M first
	mut mp := gsched.midle
	if mp != unsafe { nil } {
		gsched.midle = mp.sched_link
		gsched.nmidle--
		gsched.mu.release()
		// Give it the P and wake it
		mut p := unsafe { pp }
		wire_p(mut mp, mut p)
		mp.park.post()
		return
	}

	// No idle M available - create a new one
	id := gsched.mnext
	gsched.mnext++
	gsched.mu.release()

	new_m(id, pp)
}

// new_m creates a new OS thread (M) and associates it with a P.
fn new_m(id i64, pp &Processor) {
	mut mp := &Machine{
		id: id
		g0: &Goroutine{
			status: .running
		}
	}
	mut p := unsafe { pp }
	wire_p(mut mp, mut p)
	mp.thread = spawn m_thread_entry(mut mp)
}

// m_thread_entry is the entry point for new M (OS thread) goroutine scheduling loops.
fn m_thread_entry(mut mp Machine) {
	// Register this M as the current thread's Machine so that
	// get_current_m()/get_current_p() work on worker threads.
	set_current_m(mp)
	// Enter the scheduling loop - this never returns
	schedule_loop(mut mp)
}

// schedule_loop is the main loop for an M. It repeatedly finds and runs goroutines.
fn schedule_loop(mut mp Machine) {
	for {
		if gsched.stopped {
			return
		}
		mut pp := mp.p
		if pp == unsafe { nil } {
			acquire_p(mut mp)
			pp = mp.p
			if pp == unsafe { nil } {
				// No P available, park
				park_m(mut mp)
				continue
			}
		}
		pp.sched_tick++

		gp, _ := find_runnable(mut mp, mut pp)
		if gp == unsafe { nil } {
			// No work - try spinning briefly before parking
			if !mp.spinning {
				mp.spinning = true
				C.goroutines_atomic_fetch_add_i32(&gsched.nmspinning, 1)
			}
			// Spin a bit
			mut found := false
			for _ in 0 .. 20 {
				gp2, _ := find_runnable(mut mp, mut pp)
				if gp2 != unsafe { nil } {
					mp.spinning = false
					C.goroutines_atomic_fetch_sub_i32(&gsched.nmspinning, 1)
					execute(mut mp, mut pp, gp2)
					found = true
					break
				}
				// Yield to avoid burning CPU
				proc_yield(10)
			}
			if !found {
				if mp.spinning {
					mp.spinning = false
					C.goroutines_atomic_fetch_sub_i32(&gsched.nmspinning, 1)
				}
				park_m(mut mp)
			}
			continue
		}

		if mp.spinning {
			mp.spinning = false
			C.goroutines_atomic_fetch_sub_i32(&gsched.nmspinning, 1)
		}

		execute(mut mp, mut pp, gp)
	}
}

// steal_work tries to steal goroutines from other P's run queues.
// Translated from Go's stealWork() in proc.go.
fn steal_work(mut thisp Processor) &Goroutine {
	n := gsched.allp.len
	if n <= 1 {
		return unsafe { nil }
	}
	// Randomize starting point to avoid contention
	start := u32(C.rand()) % u32(n)
	for i := u32(0); i < u32(n); i++ {
		idx := (start + i) % u32(n)
		pp := gsched.allp[idx]
		if pp == thisp {
			continue
		}
		// Try to steal half of the target's run queue
		mut target := unsafe { pp }
		gp := runq_steal(mut target, mut thisp)
		if gp != unsafe { nil } {
			return gp
		}
	}
	return unsafe { nil }
}

// runq_steal steals half of pp's local run queue.
// Translated from Go's runqgrab/runqsteal in proc.go.
fn runq_steal(mut pp Processor, mut thisp Processor) &Goroutine {
	t := C.goroutines_atomic_load_u32(&pp.runq_tail)
	h := C.goroutines_atomic_load_u32(&pp.runq_head)
	n := t - h
	if n == 0 {
		// Try runnext
		next := pp.runnext
		if next != unsafe { nil } {
			if C.goroutines_atomic_cas_ptr(voidptr(&pp.runnext), voidptr(&next), unsafe { nil }) {
				return next
			}
		}
		return unsafe { nil }
	}
	// Steal half
	steal := n - n / 2
	mut first := unsafe { &Goroutine(nil) }
	for i := u32(0); i < steal; i++ {
		mut gp := pp.runq[(h + i) % local_queue_size]
		if i == 0 {
			first = gp
		} else {
			// Enqueue remaining stolen goroutines into thief's local run queue
			runq_put(mut thisp, gp, false)
		}
	}
	C.goroutines_atomic_fetch_add_u32(&pp.runq_head, steal)
	return first
}

// Global run queue operations (translated from Go's globrunqput/get)
fn glob_runq_put(gp &Goroutine) {
	gsched.mu.acquire()
	gsched.runq.push_back(gp)
	gsched.mu.release()
}

fn glob_runq_get() &Goroutine {
	gsched.mu.acquire()
	gp := gsched.runq.pop()
	gsched.mu.release()
	return gp
}

// Local run queue operations (translated from Go's runqput/runqget)

// runq_put puts a G on the local run queue.
// If next is true, it goes into runnext for immediate scheduling.
// Translated from Go's runqput() in proc.go.
fn runq_put(mut pp Processor, gp &Goroutine, next bool) {
	if next {
		// Fast path: put as runnext
		old := pp.runnext
		pp.runnext = unsafe { gp }
		if old == unsafe { nil } {
			return
		}
		// Kick old runnext to the regular queue
		runq_put(mut pp, old, false)
		return
	}

	// Regular path: put on the ring buffer
	h := C.goroutines_atomic_load_u32(&pp.runq_head)
	t := pp.runq_tail
	if t - h < local_queue_size {
		pp.runq[t % local_queue_size] = unsafe { gp }
		C.goroutines_atomic_store_u32(&pp.runq_tail, t + 1)
		return
	}
	// Queue is full - put half on global queue
	runq_put_slow(mut pp, gp, h, t)
}

// runq_put_slow moves half the local queue to the global queue.
// Translated from Go's runqputslow() in proc.go.
fn runq_put_slow(mut pp Processor, gp &Goroutine, h u32, t u32) {
	n := (t - h) / 2
	mut batch := GoroutineQueue{}
	for i := u32(0); i < n; i++ {
		g := pp.runq[(h + i) % local_queue_size]
		batch.push_back(g)
	}
	C.goroutines_atomic_fetch_add_u32(&pp.runq_head, n)
	batch.push_back(gp)

	gsched.mu.acquire()
	for !batch.empty() {
		gsched.runq.push_back(batch.pop())
	}
	gsched.mu.release()
}

// runq_get gets a G from the local run queue.
// Translated from Go's runqget() in proc.go.
fn runq_get(mut pp Processor) (&Goroutine, bool) {
	// Check runnext first (fast path)
	next := pp.runnext
	if next != unsafe { nil } {
		pp.runnext = unsafe { nil }
		return next, true
	}

	// Regular queue
	for {
		h := C.goroutines_atomic_load_u32(&pp.runq_head)
		t := pp.runq_tail
		if t == h {
			return unsafe { nil }, false
		}
		gp := pp.runq[h % local_queue_size]
		if C.goroutines_atomic_cas_u32(&pp.runq_head, &h, h + 1) {
			return gp, false
		}
	}
	return unsafe { nil }, false
}

// Idle P list operations
fn pid_get() &Processor {
	pp := gsched.pidle
	if pp != unsafe { nil } {
		gsched.pidle = pp.link
		gsched.npidle--
		mut p := unsafe { pp }
		p.status = .running
		p.link = unsafe { nil }
	}
	return pp
}

fn pid_put(pp &Processor) {
	mut p := unsafe { pp }
	p.status = .idle
	p.m = unsafe { nil }
	p.link = gsched.pidle
	gsched.pidle = p
	gsched.npidle++
}

// G free list operations (translated from Go's gfput/gfget)
fn gfput(mut pp Processor, gp &Goroutine) {
	mut g := unsafe { gp }
	g.status = .dead
	g.fn_ptr = unsafe { nil }
	g.fn_arg = unsafe { nil }
	g.sched_link = unsafe { nil }
	pp.g_free.push(g)
}

fn gfget(mut pp Processor) &Goroutine {
	return pp.g_free.pop()
}

fn gfput_global(gp &Goroutine) {
	gsched.g_free_mu.acquire()
	gsched.g_free.push(unsafe { gp })
	gsched.g_free_count++
	gsched.g_free_mu.release()
}

fn gfget_global() &Goroutine {
	gsched.g_free_mu.acquire()
	gp := gsched.g_free.pop()
	if gp != unsafe { nil } {
		gsched.g_free_count--
	}
	gsched.g_free_mu.release()
	return gp
}

// assign_goid allocates a unique goroutine ID.
// Uses per-P caching to avoid contention (like Go's goidcache).
fn assign_goid() u64 {
	mut pp := get_current_p()
	if pp != unsafe { nil } && pp.goid_cache < pp.goid_cache_end {
		id := pp.goid_cache
		unsafe {
			pp.goid_cache++
		}
		return id
	}
	// Refill cache from global counter
	batch := u64(16)
	id := C.goroutines_atomic_fetch_add_u64(&gsched.goid_gen, batch)
	if pp != unsafe { nil } {
		unsafe {
			pp.goid_cache = id + 1
			pp.goid_cache_end = id + batch
		}
	}
	return id
}

// proc_yield spins for a short time (used during work stealing).
fn proc_yield(count int) {
	for _ in 0 .. count {
		// CPU pause instruction to reduce power and contention
		$if amd64 {
			asm volatile amd64 {
				pause
			}
		}
		$if arm64 {
			asm volatile arm64 {
				yield
			}
		}
	}
}

// get_current_m returns the M for the current OS thread.
// Uses thread-local storage via C _Thread_local (see tls.c).
fn get_current_m() &Machine {
	return unsafe { &Machine(C.goroutines_get_current_m()) }
}

fn set_current_m(mp &Machine) {
	C.goroutines_set_current_m(voidptr(mp))
}

// get_current_p returns the P for the current OS thread's M.
fn get_current_p() &Processor {
	mp := get_current_m()
	if mp == unsafe { nil } {
		return unsafe { nil }
	}
	return mp.p
}

// get_current_g returns the currently running G.
pub fn get_current_g() &Goroutine {
	mp := get_current_m()
	if mp == unsafe { nil } {
		return unsafe { nil }
	}
	return mp.curg
}