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

import v2.types

pub struct TargetData {
pub:
	ptr_size      int
	endian_little bool
}

@[heap]
pub struct Module {
pub mut:
	name       string
	target     TargetData
	type_store TypeStore
	// For parallel SSA workers: shared reference to main module's type_store.
	// When non-nil, type operations use this instead of the local type_store.
	shared_type_store &TypeStore = unsafe { nil }

	// Type checker environment (optional, for backends that need rich type info)
	env &types.Environment = unsafe { nil }

	// Arenas
	values  []Value
	instrs  []Instruction
	blocks  []BasicBlock
	funcs   []Function
	globals []GlobalVar

	// C struct names: TypeID -> C struct name (e.g., 114 -> "stat" for struct stat)
	// Used by the C gen to emit `typedef struct <name> Struct_N;` instead of
	// generating a custom struct definition that would have the wrong memory layout.
	c_struct_names map[int]string
	// C structs marked with @[typedef] – already a C typedef, not a struct tag.
	c_typedef_structs map[int]bool

	// Constant cache: (type, name) -> ValueID for deduplication
	const_cache map[string]ValueID
}

pub fn Module.new(name string) &Module {
	mut m := &Module{
		name:       name
		type_store: TypeStore.new()
	}
	// Pre-allocate arenas to avoid repeated reallocation during SSA building.
	// A typical hello.v compilation needs ~158K values, ~134K instrs, ~10K blocks, ~1.8K funcs.
	// Pre-allocating avoids ARM64 backend issues with array growth reallocation.
	m.values = []Value{cap: 262144}
	m.instrs = []Instruction{cap: 262144}
	m.blocks = []BasicBlock{cap: 16384}
	m.funcs = []Function{cap: 4096}
	m.globals = []GlobalVar{cap: 2048}
	// Reserve ID 0 to represent "null" or "invalid", avoiding collisions
	// with map lookups returning 0.
	m.values << Value{
		kind: .unknown
		id:   0
	}
	return m
}

// release_outer_arenas_after_mir_lower releases SSA arena buffers that MIR has
// copied by value. Nested slices (value uses, instruction operands, block edges,
// function params/blocks) are intentionally not freed because MIR shallow-shares
// them after lowering.
pub fn (mut m Module) release_outer_arenas_after_mir_lower() {
	unsafe {
		m.values.free()
		m.instrs.free()
		m.blocks.free()
		m.funcs.free()
	}
	m.values = []Value{}
	m.instrs = []Instruction{}
	m.blocks = []BasicBlock{}
	m.funcs = []Function{}
}

pub fn (mut m Module) new_function(name string, ret TypeID, params []TypeID) int {
	// Check if function already exists (avoid duplicates from multiple files)
	for i, f in m.funcs {
		if f.name == name {
			return i
		}
	}
	id := m.funcs.len
	m.funcs << Function{
		id:   id
		name: name
		typ:  ret
	}
	return id
}

pub fn (mut m Module) add_block(func_id int, name string) BlockID {
	id := m.blocks.len
	// FIX: Sanitize block names for C labels (replace . with _)
	safe_name := name.replace('.', '_')
	unique_name := '${safe_name}_${id}'

	// Store 'id' (index in blocks arena) in the Value
	val_id := m.add_value_node(.basic_block, 0, unique_name, id)

	m.blocks << BasicBlock{
		id:     id
		val_id: val_id
		name:   unique_name
		parent: func_id
	}
	// Avoid m.funcs[func_id].blocks << id -- chained broken in ARM64 self-hosted
	mut f := m.funcs[func_id]
	f.blocks << id
	f.is_prototype = false
	m.funcs[func_id] = f
	return id
}

// Updated to accept 'index'
pub fn (mut m Module) add_value_node(kind ValueKind, typ TypeID, name string, index int) ValueID {
	id := m.values.len
	m.values << Value{
		id:    id
		kind:  kind
		typ:   typ
		name:  name
		index: index
	}
	return id
}

// --- Helpers to avoid chained struct-array mutations (broken in ARM64 self-hosted) ---

pub fn (mut m Module) func_add_param(func_id int, param_val ValueID) {
	mut f := m.funcs[func_id]
	f.params << param_val
	m.funcs[func_id] = f
}

pub fn (mut m Module) func_set_c_extern(func_id int, val bool) {
	mut f := m.funcs[func_id]
	f.is_c_extern = val
	m.funcs[func_id] = f
}

pub fn (mut m Module) func_set_prototype(func_id int, val bool) {
	mut f := m.funcs[func_id]
	f.is_prototype = val
	m.funcs[func_id] = f
}

pub fn (mut m Module) func_clear_blocks(func_id int) {
	mut f := m.funcs[func_id]
	f.blocks.clear()
	m.funcs[func_id] = f
}

pub fn (mut m Module) func_clear_params(func_id int) {
	mut f := m.funcs[func_id]
	f.params.clear()
	m.funcs[func_id] = f
}

pub fn (mut m Module) func_set_params(func_id int, params []ValueID) {
	mut f := m.funcs[func_id]
	f.params = params
	m.funcs[func_id] = f
}

pub fn (mut m Module) block_add_succ(from int, to int) {
	mut blk := m.blocks[from]
	blk.succs << to
	m.blocks[from] = blk
}

pub fn (mut m Module) block_add_pred(to int, from int) {
	mut blk := m.blocks[to]
	blk.preds << from
	m.blocks[to] = blk
}

pub fn (mut m Module) nop_instr(val_idx int) {
	mut instr := m.instrs[val_idx]
	instr.op = .bitcast
	instr.operands = []
	m.instrs[val_idx] = instr
}

// Get or create a constant value, reusing existing ones when possible.
// This maintains SSA's immutability principle by avoiding duplicate constants.
pub fn (mut m Module) get_or_add_const(typ TypeID, name string) ValueID {
	key := '${typ}:${name}'
	if existing := map_get_value_id(m.const_cache, key) {
		return existing
	}
	id := m.add_value_node(.constant, typ, name, 0)
	m.const_cache[key] = id
	return id
}

pub fn (m Module) get_block_from_val(val_id int) int {
	return m.values[val_id].index
}

pub fn (mut m Module) add_instr(op OpCode, block BlockID, typ TypeID, operands []ValueID) ValueID {
	// 1. Save Instruction Index
	instr_idx := m.instrs.len

	instr := Instruction{
		op:       op
		block:    block
		typ:      typ
		operands: operands
	}
	m.instrs << instr

	// 2. Pass instr_idx to Value
	val_id := m.add_value_node(.instruction, typ, '', instr_idx)

	// 3. Link Block — read whole struct, modify, write back (chained broken in ARM64)
	mut blk := m.blocks[block]
	blk.instrs << val_id
	m.blocks[block] = blk

	// 4. Update Def-Use — same pattern
	for op_id in operands {
		if op_id < m.values.len {
			mut v := m.values[op_id]
			v.uses << val_id
			m.values[op_id] = v
		}
	}

	return val_id
}

pub fn (mut m Module) add_global(name string, typ TypeID, is_const bool) int {
	return m.add_global_with_value(name, typ, is_const, 0)
}

pub fn (mut m Module) add_global_with_value(name string, typ TypeID, is_const bool, initial_value i64) int {
	id := m.globals.len
	g := GlobalVar{
		name:          name
		typ:           typ
		linkage:       .private // Local global, not external
		is_constant:   is_const
		initial_value: initial_value
	}
	m.globals << g

	// FIX: The Value representing a global is a POINTER to the data
	ptr_typ := m.type_store.get_ptr(typ)
	return m.add_value_node(.global, ptr_typ, name, id)
}

pub fn (mut m Module) add_global_with_data(name string, elem_type TypeID, is_const bool, data []u8) int {
	id := m.globals.len
	g := GlobalVar{
		name:         name
		typ:          elem_type
		linkage:      .private
		is_constant:  is_const
		initial_data: data
	}
	m.globals << g

	// The Value is a POINTER to element data (for direct indexing)
	ptr_typ := m.type_store.get_ptr(elem_type)
	return m.add_value_node(.global, ptr_typ, name, id)
}

// add_external_global adds an external global variable (defined outside this module)
// Returns the ValueID for the global pointer
pub fn (mut m Module) add_external_global(name string, typ TypeID) ValueID {
	// Check if this external global already exists
	for v in m.values {
		if v.kind == .global && v.name == name {
			return v.id
		}
	}

	// Create a new external global
	id := m.globals.len
	g := GlobalVar{
		name:    name
		typ:     typ
		linkage: .external
	}
	m.globals << g

	// The Value representing a global is a POINTER to the data
	ptr_typ := m.type_store.get_ptr(typ)
	return m.add_value_node(.global, ptr_typ, name, id)
}

pub fn (mut m Module) add_instr_front(op OpCode, block BlockID, typ TypeID, operands []ValueID) ValueID {
	instr_idx := m.instrs.len
	instr := Instruction{
		op:       op
		block:    block
		typ:      typ
		operands: operands
	}
	m.instrs << instr
	val_id := m.add_value_node(.instruction, typ, '', instr_idx)

	// Prepend to block instructions — read whole struct, modify, write back (chained broken in ARM64)
	mut blk := m.blocks[block]
	blk.instrs.prepend(val_id)
	m.blocks[block] = blk

	// Update Def-Use — same pattern
	for op_id in operands {
		if op_id < m.values.len {
			mut v := m.values[op_id]
			v.uses << val_id
			m.values[op_id] = v
		}
	}
	return val_id
}

// append_phi_operands appends a (val, block_val) pair to a phi instruction's operands.
// Append (val, block_val) pair to phi instruction operands.
// Avoid m.instrs[idx].operands << x — chained append broken in ARM64 self-hosted.
pub fn (mut m Module) append_phi_operands(instr_idx int, val int, block_val int) {
	// Read whole struct, modify, write back (chained broken in ARM64)
	mut instr := m.instrs[instr_idx]
	instr.operands << val
	instr.operands << block_val
	m.instrs[instr_idx] = instr
}

pub fn (mut m Module) replace_uses(old_val int, new_val int) {
	// Copy uses, because we modify instr operands which might change things
	uses := m.values[old_val].uses.clone()
	for use_id in uses {
		use_val := m.values[use_id]
		if use_val.kind == .instruction {
			// Read whole instr, modify operands, write back (chained broken in ARM64)
			mut instr := m.instrs[use_val.index]
			mut replaced := false
			for i in 0 .. instr.operands.len {
				if instr.operands[i] == old_val {
					instr.operands[i] = new_val
					replaced = true
				}
			}
			if replaced {
				m.instrs[use_val.index] = instr
			}
			// Only add to uses list once per user, even if used multiple times
			if replaced && use_id !in m.values[new_val].uses {
				mut nv := m.values[new_val]
				nv.uses << use_id
				m.values[new_val] = nv
			}
		}
	}
	mut ov := m.values[old_val]
	ov.uses = []
	m.values[old_val] = ov
}

fn dfs(mut m Module, blk int, mut visited map[int]bool, mut rpo []int) {
	visited[blk] = true
	for s in m.blocks[blk].succs {
		if !visited[s] {
			dfs(mut m, s, mut visited, mut rpo)
		}
	}
	rpo << blk
}

fn (mut m Module) get_rpo(func Function) []int {
	mut visited := map[int]bool{}
	mut rpo := []int{}
	dfs(mut m, func.blocks[0], mut visited, mut rpo)
	// rpo.reverse_inplace()
	rpo = rpo.reverse()
	return rpo
}

// new_worker_module creates a lightweight Module for parallel SSA building.
// Each worker gets its own copy of type_store (COW from main).
// Has its own values/instrs/blocks arenas starting from empty.
pub fn (mut m Module) new_worker_module() &Module {
	// Explicitly clone funcs and globals to avoid COW races between threads.
	mut wf := []Function{cap: m.funcs.len}
	for f in m.funcs {
		wf << f
	}
	mut wg := []GlobalVar{cap: m.globals.len}
	for g in m.globals {
		wg << g
	}
	// Deep-clone TypeStore to avoid COW races on types[] and cache map.
	mut wts := TypeStore{}
	wts.types = []Type{cap: m.type_store.types.len}
	for t in m.type_store.types {
		wts.types << t
	}
	wts.cache = m.type_store.cache.clone()

	mut w := &Module{
		name:              m.name
		shared_type_store: unsafe { nil }
		type_store:        wts
		env:               m.env
		funcs:             wf
		globals:           wg
	}
	// Seed worker with main module's values so that ValueIDs from Phases 1-3
	// (constants, globals, func_refs, params) are valid in the worker.
	w.values = []Value{cap: m.values.len + 32768}
	for v in m.values {
		w.values << v
	}
	// Also seed instrs and blocks (typically empty after Phases 1-3, but just in case).
	w.instrs = []Instruction{cap: m.instrs.len + 32768}
	for instr in m.instrs {
		w.instrs << instr
	}
	w.blocks = []BasicBlock{cap: m.blocks.len + 2048}
	for blk in m.blocks {
		w.blocks << blk
	}
	return w
}

// FuncSSAData holds the SSA data produced by a worker for a single function.
pub struct FuncSSAData {
pub:
	func_idx int       // Index into main module's funcs[]
	blocks   []BlockID // Worker-local block IDs
	params   []ValueID // Worker-local param ValueIDs
}

// merge_worker_module merges a worker's SSA arenas into the main module.
// Workers are seeded with the main module's values/instrs/blocks from Phases 1-3.
// seed_values/seed_instrs/seed_blocks/seed_funcs are the lengths of the seed data.
// Only data beyond the seed is new and needs to be merged with ID remapping.
// func_data contains (func_idx, blocks, params) for updating main funcs[].
pub fn (mut m Module) merge_worker_module(w &Module, func_data []FuncSSAData, seed_values int, seed_instrs int, seed_blocks int, seed_types int, seed_funcs int) {
	// Build type remapping: worker type IDs >= seed_types may differ from main.
	// For each new worker type, find or create equivalent in main.
	mut type_remap := []TypeID{len: w.type_store.types.len, init: 0}
	// Seed types map to themselves (identity)
	for ti := 0; ti < seed_types && ti < type_remap.len; ti++ {
		type_remap[ti] = ti
	}
	// Map new worker types to main types
	for ti := seed_types; ti < w.type_store.types.len; ti++ {
		wt := w.type_store.types[ti]
		// Generate cache key matching TypeStore conventions
		mut cache_key := ''
		match wt.kind {
			.int_t {
				cache_key = if wt.is_unsigned { 'u${wt.width}' } else { 'i${wt.width}' }
			}
			.float_t {
				cache_key = 'f${wt.width}'
			}
			.ptr_t {
				// Remap the elem_type first
				remapped_elem := if int(wt.elem_type) < type_remap.len {
					type_remap[int(wt.elem_type)]
				} else {
					wt.elem_type
				}
				cache_key = 'p${remapped_elem}'
			}
			.array_t {
				remapped_elem := if int(wt.elem_type) < type_remap.len {
					type_remap[int(wt.elem_type)]
				} else {
					wt.elem_type
				}
				cache_key = 'a${remapped_elem}_${wt.len}'
			}
			else {}
		}

		// Try cache lookup in main
		if cache_key.len > 0 {
			if existing := map_get_type_id(m.type_store.cache, cache_key) {
				type_remap[ti] = existing
				continue
			}
		}
		// Register new type in main with remapped field types
		remapped_ret := if int(wt.ret_type) < type_remap.len && int(wt.ret_type) >= seed_types {
			type_remap[int(wt.ret_type)]
		} else {
			wt.ret_type
		}
		remapped_elem2 := if int(wt.elem_type) < type_remap.len && int(wt.elem_type) >= seed_types {
			type_remap[int(wt.elem_type)]
		} else {
			wt.elem_type
		}
		mut new_type := Type{
			kind:        wt.kind
			width:       wt.width
			len:         wt.len
			is_c_struct: wt.is_c_struct
			is_union:    wt.is_union
			is_unsigned: wt.is_unsigned
			ret_type:    remapped_ret
			elem_type:   remapped_elem2
		}
		if wt.fields.len > 0 {
			mut new_fields := []TypeID{cap: wt.fields.len}
			for f in wt.fields {
				new_fields << if int(f) < type_remap.len && int(f) >= seed_types {
					type_remap[int(f)]
				} else {
					f
				}
			}
			new_type = Type{
				...new_type
				fields:      new_fields
				field_names: wt.field_names
			}
		}
		if wt.params.len > 0 {
			mut new_params := []TypeID{cap: wt.params.len}
			for p in wt.params {
				new_params << if int(p) < type_remap.len && int(p) >= seed_types {
					type_remap[int(p)]
				} else {
					p
				}
			}
			new_type = Type{
				...new_type
				params: new_params
			}
		}
		new_id := m.type_store.register(new_type)
		if cache_key.len > 0 {
			m.type_store.cache[cache_key] = new_id
		}
		type_remap[ti] = new_id
	}

	// Remapping offsets: worker IDs in [seed..seed+N) → main IDs in [main.len..main.len+N).
	// For seed IDs (< seed_values), no remapping needed — they're the same in both.
	value_off := m.values.len - seed_values
	instr_off := m.instrs.len - seed_instrs
	block_off := m.blocks.len - seed_blocks

	// Append worker values beyond seed
	for wi := seed_values; wi < w.values.len; wi++ {
		wv := w.values[wi]
		mut new_index := wv.index
		if wv.kind == .instruction {
			new_index += instr_off
		} else if wv.kind == .basic_block {
			new_index += block_off
		}
		// Remap uses — only IDs >= seed need remapping
		mut new_uses := []ValueID{cap: wv.uses.len}
		for u in wv.uses {
			new_uses << if u >= seed_values { u + value_off } else { u }
		}
		// Remap type
		new_typ := if int(wv.typ) >= seed_types && int(wv.typ) < type_remap.len {
			type_remap[int(wv.typ)]
		} else {
			wv.typ
		}
		m.values << Value{
			id:    wi + value_off
			kind:  wv.kind
			typ:   new_typ
			name:  wv.name
			index: new_index
			uses:  new_uses
		}
	}

	// Append worker instructions beyond seed
	for ii := seed_instrs; ii < w.instrs.len; ii++ {
		instr := w.instrs[ii]
		mut new_ops := []ValueID{cap: instr.operands.len}
		for op in instr.operands {
			new_ops << if op >= seed_values { op + value_off } else { op }
		}
		// Remap type
		new_typ := if int(instr.typ) >= seed_types && int(instr.typ) < type_remap.len {
			type_remap[int(instr.typ)]
		} else {
			instr.typ
		}
		m.instrs << Instruction{
			op:         instr.op
			block:      if instr.block >= seed_blocks {
				instr.block + block_off
			} else {
				instr.block
			}
			typ:        new_typ
			operands:   new_ops
			pos:        instr.pos
			atomic_ord: instr.atomic_ord
			inline:     instr.inline
		}
	}

	// Append worker blocks beyond seed
	for bi := seed_blocks; bi < w.blocks.len; bi++ {
		blk := w.blocks[bi]
		mut new_instrs := []ValueID{cap: blk.instrs.len}
		for vi in blk.instrs {
			new_instrs << if vi >= seed_values { vi + value_off } else { vi }
		}
		mut new_preds := []BlockID{cap: blk.preds.len}
		for p in blk.preds {
			new_preds << if p >= seed_blocks { p + block_off } else { p }
		}
		mut new_succs := []BlockID{cap: blk.succs.len}
		for s in blk.succs {
			new_succs << if s >= seed_blocks { s + block_off } else { s }
		}
		m.blocks << BasicBlock{
			id:     blk.id + block_off
			val_id: if blk.val_id >= seed_values { blk.val_id + value_off } else { blk.val_id }
			name:   blk.name
			parent: blk.parent
			instrs: new_instrs
			preds:  new_preds
			succs:  new_succs
		}
	}

	// Update main module's funcs[] with blocks and params from worker
	for fd in func_data {
		if fd.func_idx >= m.funcs.len {
			continue
		}
		mut f := m.funcs[fd.func_idx]
		// Skip if function was already built by a previous worker's merge.
		// Without this guard, duplicate blocks from multiple workers get appended,
		// leading to mismatched parameter/value IDs and crashes.
		if f.blocks.len > 0 {
			continue
		}
		for blk_id in fd.blocks {
			f.blocks << if blk_id >= seed_blocks { blk_id + block_off } else { blk_id }
		}
		mut new_params := []ValueID{cap: fd.params.len}
		for p in fd.params {
			new_params << if p >= seed_values { p + value_off } else { p }
		}
		f.params = new_params
		m.funcs[fd.func_idx] = f
	}

	// Also add any new functions created by workers (stubs like wyhash, array_eq).
	for fi := seed_funcs; fi < w.funcs.len; fi++ {
		wfunc := w.funcs[fi]
		mut already_added := false
		for existing in m.funcs {
			if existing.name == wfunc.name {
				already_added = true
				break
			}
		}
		if already_added {
			continue
		}
		mut new_blocks := []BlockID{cap: wfunc.blocks.len}
		for blk_id in wfunc.blocks {
			new_blocks << if blk_id >= seed_blocks { blk_id + block_off } else { blk_id }
		}
		mut new_params := []ValueID{cap: wfunc.params.len}
		for p in wfunc.params {
			new_params << if p >= seed_values { p + value_off } else { p }
		}
		// Remap function return type
		new_typ := if int(wfunc.typ) >= seed_types && int(wfunc.typ) < type_remap.len {
			type_remap[int(wfunc.typ)]
		} else {
			wfunc.typ
		}
		m.funcs << Function{
			id:     m.funcs.len
			name:   wfunc.name
			typ:    new_typ
			blocks: new_blocks
			params: new_params
		}
	}
}