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

import v2.ssa
import v2.types

pub enum ValueKind {
	unknown
	constant
	argument
	global
	instruction
	basic_block
	string_literal
	c_string_literal
	func_ref
}

// str returns the symbolic name for a MIR value kind.
pub fn (k ValueKind) str() string {
	return match k {
		.unknown { 'unknown' }
		.constant { 'constant' }
		.argument { 'argument' }
		.global { 'global' }
		.instruction { 'instruction' }
		.basic_block { 'basic_block' }
		.string_literal { 'string_literal' }
		.c_string_literal { 'c_string_literal' }
		.func_ref { 'func_ref' }
	}
}

pub enum AbiArgClass as u8 {
	in_reg
	indirect
}

// Classification metadata for future ABI-aware codegen.  The active x64
// lowering path still consumes the legacy abi_*_class and abi_ret_indirect
// projection fields in this palier.
pub enum AbiPassMode {
	direct
	indirect
}

pub enum AbiEightbyteClass {
	no_class
	integer
	sse
	sseup
	memory
}

pub struct AbiValueClass {
pub mut:
	mode    AbiPassMode
	size    int
	classes []AbiEightbyteClass
}

pub enum AbiLocationKind {
	none
	int_reg
	sse_reg
	stack
}

pub struct AbiLocation {
pub mut:
	kind   AbiLocationKind
	index  int
	offset int
	class  AbiEightbyteClass
}

pub struct AbiValueLayout {
pub mut:
	value_class AbiValueClass
	locs        []AbiLocation
}

pub struct Value {
pub:
	id    int
	typ   ssa.TypeID
	index int
pub mut:
	kind ValueKind
	name string
	uses []ssa.ValueID
}

pub struct Instruction {
pub mut:
	op               ssa.OpCode
	operands         []ssa.ValueID
	abi_ret_indirect bool
	abi_arg_class    []AbiArgClass
	abi_ret_class    AbiValueClass
	abi_arg_classes  []AbiValueClass
	abi_arg_layouts  []AbiValueLayout
pub:
	typ       ssa.TypeID
	block     int
	src_index int
}

pub struct BasicBlock {
pub:
	id     int
	val_id int
	name   string
	parent int
pub mut:
	instrs []ssa.ValueID
	preds  []ssa.BlockID
	succs  []ssa.BlockID
}

pub struct Function {
pub:
	id          int
	name        string
	typ         ssa.TypeID
	linkage     ssa.Linkage
	call_conv   ssa.CallConv
	is_c_extern bool // C-language extern function (provided by libc/system libraries)
pub mut:
	blocks            []ssa.BlockID
	params            []ssa.ValueID
	abi_ret_indirect  bool
	abi_param_class   []AbiArgClass
	abi_ret_class     AbiValueClass
	abi_param_classes []AbiValueClass
	abi_param_layouts []AbiValueLayout
}

@[heap]
pub struct Module {
pub:
	name    string
	target  ssa.TargetData
	ssa_mod &ssa.Module        = unsafe { nil }
	env     &types.Environment = unsafe { nil }
pub mut:
	type_store ssa.TypeStore
	values     []Value
	instrs     []Instruction
	blocks     []BasicBlock
	funcs      []Function
	globals    []ssa.GlobalVar
}

fn share_value_ids(values []ssa.ValueID) []ssa.ValueID {
	return values
}

fn share_block_ids(blocks []ssa.BlockID) []ssa.BlockID {
	return blocks
}

pub fn lower_from_ssa(ssa_mod &ssa.Module) Module {
	mut mod := Module{
		name:       ssa_mod.name
		target:     ssa_mod.target
		ssa_mod:    ssa_mod
		env:        ssa_mod.env
		type_store: ssa_mod.type_store
		values:     []Value{len: ssa_mod.values.len}
		instrs:     []Instruction{len: ssa_mod.instrs.len}
		blocks:     []BasicBlock{len: ssa_mod.blocks.len}
		funcs:      []Function{len: ssa_mod.funcs.len}
		globals:    ssa_mod.globals
	}

	for i, val in ssa_mod.values {
		mod.values[i] = Value{
			id:    val.id
			typ:   val.typ
			index: val.index
			kind:  value_kind_from_ssa(val.kind)
			name:  val.name
			uses:  share_value_ids(val.uses)
		}
	}

	for i, instr in ssa_mod.instrs {
		mod.instrs[i] = Instruction{
			op:               instr.op
			operands:         share_value_ids(instr.operands)
			abi_ret_indirect: false
			abi_arg_class:    []AbiArgClass{}
			abi_ret_class:    AbiValueClass{}
			abi_arg_classes:  []AbiValueClass{}
			abi_arg_layouts:  []AbiValueLayout{}
			typ:              instr.typ
			block:            instr.block
			src_index:        i
		}
	}

	for i, blk in ssa_mod.blocks {
		mod.blocks[i] = BasicBlock{
			id:     blk.id
			val_id: blk.val_id
			name:   blk.name
			parent: blk.parent
			instrs: share_value_ids(blk.instrs)
			preds:  share_block_ids(blk.preds)
			succs:  share_block_ids(blk.succs)
		}
	}

	for i, f in ssa_mod.funcs {
		mod.funcs[i] = Function{
			id:                f.id
			name:              f.name
			typ:               f.typ
			linkage:           f.linkage
			call_conv:         f.call_conv
			is_c_extern:       f.is_c_extern
			blocks:            share_block_ids(f.blocks)
			params:            share_value_ids(f.params)
			abi_ret_indirect:  false
			abi_param_class:   []AbiArgClass{len: f.params.len, init: .in_reg}
			abi_ret_class:     AbiValueClass{}
			abi_param_classes: []AbiValueClass{len: f.params.len}
			abi_param_layouts: []AbiValueLayout{len: f.params.len}
		}
	}

	return mod
}

// release_after_native_codegen releases MIR storage after native codegen has
// copied everything it needs into the object writer/linker.
pub fn (mut m Module) release_after_native_codegen() {
	unsafe {
		m.values.free()
		m.instrs.free()
		m.blocks.free()
		m.funcs.free()
		m.globals.free()
		m.type_store.types.free()
		m.type_store.cache.free()
	}
	m.values = []Value{}
	m.instrs = []Instruction{}
	m.blocks = []BasicBlock{}
	m.funcs = []Function{}
	m.globals = []ssa.GlobalVar{}
	m.type_store = ssa.TypeStore{}
}

pub fn (m &Module) ssa() &ssa.Module {
	return m.ssa_mod
}

pub fn (m &Module) type_size(typ_id ssa.TypeID) int {
	mut visiting := map[ssa.TypeID]bool{}
	return m.type_size_with_seen(typ_id, mut visiting)
}

fn (m &Module) type_size_with_seen(typ_id ssa.TypeID, mut visiting map[ssa.TypeID]bool) int {
	if m.ssa_mod == unsafe { nil } || typ_id < 0 || typ_id >= m.ssa_mod.type_store.types.len {
		return 0
	}
	if typ_id == 0 {
		return 0
	}
	if visiting[typ_id] {
		return 8
	}
	visiting[typ_id] = true
	typ := m.ssa_mod.type_store.types[typ_id]
	size := match typ.kind {
		.void_t {
			0
		}
		.int_t {
			if typ.width > 0 { (typ.width + 7) / 8 } else { 8 }
		}
		.float_t {
			if typ.width > 0 { (typ.width + 7) / 8 } else { 8 }
		}
		.ptr_t {
			8
		}
		.array_t {
			elem_size := m.type_size_with_seen(typ.elem_type, mut visiting)
			typ.len * elem_size
		}
		.struct_t {
			mut total := 0
			mut max_align := 1
			for field_typ in typ.fields {
				align := m.type_align_with_seen(field_typ, mut visiting)
				if align > max_align {
					max_align = align
				}
				if align > 1 && total % align != 0 {
					total = (total + align - 1) & ~(align - 1)
				}
				total += m.type_size_with_seen(field_typ, mut visiting)
			}
			// Align struct size to its largest field alignment
			if max_align > 1 && total % max_align != 0 {
				total = (total + max_align - 1) & ~(max_align - 1)
			}
			if total > 0 {
				total
			} else {
				8
			}
		}
		.func_t {
			8
		}
		.label_t {
			0
		}
		.metadata_t {
			0
		}
	}

	visiting.delete(typ_id)
	return size
}

pub fn (m &Module) type_align(typ_id ssa.TypeID) int {
	mut visiting := map[ssa.TypeID]bool{}
	return m.type_align_with_seen(typ_id, mut visiting)
}

fn (m &Module) type_align_with_seen(typ_id ssa.TypeID, mut visiting map[ssa.TypeID]bool) int {
	if m.ssa_mod != unsafe { nil } && typ_id > 0 && typ_id < m.ssa_mod.type_store.types.len {
		if visiting[typ_id] {
			return 8
		}
		typ := m.ssa_mod.type_store.types[typ_id]
		if typ.kind == .array_t {
			return m.type_align_with_seen(typ.elem_type, mut visiting)
		}
	}
	size := m.type_size_with_seen(typ_id, mut visiting)
	if size >= 8 {
		return 8
	}
	if size >= 4 {
		return 4
	}
	if size >= 2 {
		return 2
	}
	return 1
}

fn value_kind_from_ssa(kind ssa.ValueKind) ValueKind {
	return match kind {
		.unknown { .unknown }
		.constant { .constant }
		.argument { .argument }
		.global { .global }
		.instruction { .instruction }
		.basic_block { .basic_block }
		.string_literal { .string_literal }
		.c_string_literal { .c_string_literal }
		.func_ref { .func_ref }
	}
}