package shader

// IR data structures for the Luma shader compiler.
// This is a structured IR that preserves control flow nesting (if/for/while)
// rather than using CFG basic blocks. This makes text backend emission trivial.
// Phase 3 can lower to CFG for SPIR-V.

IR_Module :: struct {
	functions:      [dynamic]IR_Function,
	structs:        [dynamic]IR_Struct,
	bindings:       [dynamic]IR_Binding,
	spec_constants: [dynamic]IR_Spec_Constant,
	shared_vars:    [dynamic]IR_Shared_Var,
}

IR_Address_Space :: enum {
	Private,
	Uniform,
	Storage,
	Workgroup,
	Input,
	Output,
	Push_Constant,
}

IR_Shared_Var :: struct {
	name: string,
	type: ^Resolved_Type,
}

IR_Var_Id :: distinct int

IR_Var_Decl :: struct {
	id:   IR_Var_Id,
	name: string,
	type: ^Resolved_Type,
}

Barrier_Kind :: enum {
	Workgroup,
}

IR_Spec_Constant :: struct {
	name:          string,
	spec_id:       int,
	type:          ^Resolved_Type,
	default_value: IR_Const_Value,
}

IR_Const_Value :: union {
	i64,
	f64,
	bool,
}

IR_Function :: struct {
	name:           string,
	params:         []IR_Param,
	return_type:    ^Resolved_Type,
	body:           [dynamic]IR_Stmt,
	is_entry:       bool,
	stage:          Shader_Stage,
	inputs:         [dynamic]IR_IO_Var, // flattened entry point inputs
	outputs:        [dynamic]IR_IO_Var, // flattened entry point outputs
	workgroup_size: [3]int,
	var_decls:      [dynamic]IR_Var_Decl, // all local variables, indexed by IR_Var_Id
}

IR_Param :: struct {
	name: string,
	type: ^Resolved_Type,
	id:   IR_Var_Id,
}

// Flattened entry point I/O variable
IR_IO_Var :: struct {
	name:     string,
	type:     ^Resolved_Type,
	location: int,    // -1 if builtin
	builtin:  string, // "" if location-based
}

IR_Struct :: struct {
	name:   string,
	fields: []IR_Struct_Field,
}

IR_Struct_Field :: struct {
	name:       string,
	type:       ^Resolved_Type,
	attributes: []Ast_Attribute,
}

IR_Binding_Kind :: enum {
	Uniform,
	Buffer,
	Texture,
	Sampler,
	Push_Constant,
}

IR_Binding :: struct {
	name:          string,
	type:          ^Resolved_Type,
	kind:          IR_Binding_Kind,
	group:         int,
	binding_num:   int,
	struct_ref:    ^Type_Struct_Resolved, // for uniform/buffer blocks
	combined_name: string,                // original name for split sampler pairs, "" otherwise
	address_space: IR_Address_Space,
}

// Statements — structured control flow preserved

IR_Discard :: struct {
	span: Source_Span,
}

IR_Break :: struct {
	span: Source_Span,
}

IR_Continue :: struct {
	span: Source_Span,
}

IR_Barrier :: struct {
	kind: Barrier_Kind,
	span: Source_Span,
}

IR_Stmt :: union {
	^IR_Let,
	^IR_Assign,
	^IR_Return,
	^IR_If,
	^IR_For,
	^IR_While,
	^IR_Store_Output,
	^IR_Expr_Stmt,
	^IR_Barrier,
	^IR_Discard,
	^IR_Break,
	^IR_Continue,
}

IR_Let :: struct {
	name:    string,
	id:      IR_Var_Id,
	value:   ^IR_Expr,
	type:    ^Resolved_Type,
	span:    Source_Span,
	mutable: bool,
}

IR_Assign :: struct {
	target: ^IR_Expr,
	value:  ^IR_Expr,
	span:   Source_Span,
}

IR_Return :: struct {
	value: ^IR_Expr, // nil for void return
	span:  Source_Span,
}

IR_If :: struct {
	condition:      ^IR_Expr,
	then_body:      [dynamic]IR_Stmt,
	elseif_clauses: []IR_Elseif,
	else_body:      [dynamic]IR_Stmt,
	span:           Source_Span,
}

IR_Elseif :: struct {
	condition: ^IR_Expr,
	body:      [dynamic]IR_Stmt,
}

IR_For :: struct {
	var_name: string,
	var_id:   IR_Var_Id,
	start:    ^IR_Expr,
	stop:     ^IR_Expr,
	step:     ^IR_Expr, // nil if default step of 1
	body:     [dynamic]IR_Stmt,
	span:     Source_Span,
}

IR_While :: struct {
	condition: ^IR_Expr,
	body:      [dynamic]IR_Stmt,
	span:      Source_Span,
}

IR_Store_Output :: struct {
	io_index: int,     // index into IR_Function.outputs
	value:    ^IR_Expr,
	span:     Source_Span,
}

IR_Expr_Stmt :: struct {
	expr: ^IR_Expr,
	span: Source_Span,
}

// Expressions

IR_Expr :: struct {
	kind:          IR_Expr_Kind,
	type:          ^Resolved_Type,
	derived:       IR_Expr_Derived,
}

IR_Expr_Kind :: enum {
	Literal,
	Var_Ref,
	Binary,
	Unary,
	Call,
	Field_Access,
	Swizzle,
	Index,
	Construct,
	Type_Cast,
	Load_Binding,
	Input_Field,
	Builtin_Var,
	Composite_Extract, // value-semantic field access (struct.field on SSA values)
	Vector_Shuffle,    // swizzle/reorder vector components by index
	Shared_Ref,        // reference to a shared (workgroup) variable
	Select,            // ternary conditional: select(cond, true, false)
}

IR_Expr_Derived :: union {
	^IR_Literal,
	^IR_Var_Ref,
	^IR_Binary,
	^IR_Unary,
	^IR_Call,
	^IR_Field_Access,
	^IR_Swizzle,
	^IR_Index,
	^IR_Construct,
	^IR_Type_Cast,
	^IR_Load_Binding,
	^IR_Input_Field,
	^IR_Builtin_Var,
	^IR_Composite_Extract,
	^IR_Vector_Shuffle,
	^IR_Shared_Ref,
	^IR_Select,
}

IR_Select :: struct {
	condition: ^IR_Expr,
	true_val:  ^IR_Expr,
	false_val: ^IR_Expr,
}

IR_Shared_Ref :: struct {
	name: string,
}

IR_Literal :: struct {
	value: IR_Literal_Value,
}

IR_Literal_Value :: union {
	i64,
	f64,
	bool,
}

IR_Var_Ref :: struct {
	name: string,
	id:   IR_Var_Id,
}

IR_Binary :: struct {
	op:    IR_Op,
	left:  ^IR_Expr,
	right: ^IR_Expr,
}

IR_Unary :: struct {
	op:      IR_Op,
	operand: ^IR_Expr,
}

IR_Call :: struct {
	name:       string,
	args:       []^IR_Expr,
	is_builtin: bool,
}

IR_Field_Access :: struct {
	object:     ^IR_Expr,
	field_name: string,
}

IR_Swizzle :: struct {
	object:     ^IR_Expr,
	components: string, // e.g. "xyz", "rgb"
}

IR_Index :: struct {
	object: ^IR_Expr,
	index:  ^IR_Expr,
}

IR_Construct :: struct {
	type_name: string,
	args:      []^IR_Expr,
}

IR_Type_Cast :: struct {
	value: ^IR_Expr,
}

IR_Load_Binding :: struct {
	name: string,
}

IR_Input_Field :: struct {
	param_name: string,
	field_name: string,
}

IR_Builtin_Var :: struct {
	name:     string,
	stage:    Shader_Stage,
	is_input: bool,
}

// Value-semantic field extraction from a composite (struct/vector) SSA value.
// Maps to SPIR-V OpCompositeExtract. Text backends emit as obj.field.
IR_Composite_Extract :: struct {
	object:     ^IR_Expr,
	index:      int,    // field index in the composite type
	field_name: string, // field name for text backend emission
}

// Swizzle/reorder vector components by integer indices.
// Maps to SPIR-V OpVectorShuffle.
IR_Vector_Shuffle :: struct {
	object:     ^IR_Expr,
	components: []int, // e.g. [0, 1, 2] for .xyz
}

IR_Op :: enum {
	// Arithmetic
	Add, Sub, Mul, Div, Mod,
	// Comparison
	Eq, Neq, Lt, Gt, Lte, Gte,
	// Logical
	And, Or,
	// Unary
	Neg, Not,
}

// -- Debug printer --

ir_to_string :: proc(module: ^IR_Module, allocator := context.allocator) -> string {
	w := writer_init(allocator)

	write_line(&w, "=== IR Module ===")
	write_line(&w, "")

	// Structs
	for s in module.structs {
		write_line(&w, "struct ", s.name)
		indent(&w)
		for f in s.fields {
			write_line(&w, f.name, ": ", type_to_string(f.type))
		}
		dedent(&w)
		write_line(&w, "end")
		write_line(&w, "")
	}

	// Spec constants
	for sc in module.spec_constants {
		write_line(&w, "@spec(", sc.spec_id, ") const ", sc.name, ": ", type_to_string(sc.type), " = ", ir_const_value_to_string(sc.default_value))
	}
	if len(module.spec_constants) > 0 {
		write_line(&w, "")
	}

	// Shared variables
	for sv in module.shared_vars {
		write_line(&w, "shared ", sv.name, ": ", type_to_string(sv.type))
	}
	if len(module.shared_vars) > 0 {
		write_line(&w, "")
	}

	// Bindings
	for b in module.bindings {
		kind_str: string
		switch b.kind {
		case .Uniform:       kind_str = "uniform"
		case .Buffer:        kind_str = "buffer"
		case .Texture:       kind_str = "texture"
		case .Sampler:       kind_str = "sampler"
		case .Push_Constant: kind_str = "push_constant"
		}
		if b.kind == .Push_Constant {
			write_line(&w, "@push_constant uniform ", b.name, ": ", type_to_string(b.type))
		} else {
			write_line(&w, "@group(", b.group, ") @binding(", b.binding_num, ") ", kind_str, " ", b.name, ": ", type_to_string(b.type))
		}
	}
	if len(module.bindings) > 0 {
		write_line(&w, "")
	}

	// Functions
	for &fn in module.functions {
		if fn.is_entry {
			write(&w, "@entry(")
			#partial switch fn.stage {
			case .Vertex:   write(&w, "vertex")
			case .Fragment: write(&w, "fragment")
			case .Compute:  write(&w, "compute")
			case:           write(&w, "unknown")
			}
			write(&w, ") ")
		}
		write(&w, "function ", fn.name, "(")
		for p, i in fn.params {
			if i > 0 do write(&w, ", ")
			write(&w, p.name, ": ", type_to_string(p.type))
		}
		write(&w, ") -> ", type_to_string(fn.return_type), "\n")

		if fn.is_entry {
			indent(&w)
			write_line(&w, "-- inputs:")
			for io, i in fn.inputs {
				if io.builtin != "" {
					write_line(&w, "  [", i, "] @builtin(", io.builtin, ") ", io.name, ": ", type_to_string(io.type))
				} else {
					write_line(&w, "  [", i, "] @location(", io.location, ") ", io.name, ": ", type_to_string(io.type))
				}
			}
			write_line(&w, "-- outputs:")
			for io, i in fn.outputs {
				if io.builtin != "" {
					write_line(&w, "  [", i, "] @builtin(", io.builtin, ") ", io.name, ": ", type_to_string(io.type))
				} else {
					write_line(&w, "  [", i, "] @location(", io.location, ") ", io.name, ": ", type_to_string(io.type))
				}
			}
			dedent(&w)
		}

		indent(&w)
		ir_print_stmts(&w, fn.body[:])
		dedent(&w)
		write_line(&w, "end")
		write_line(&w, "")
	}

	return writer_to_string(w)
}

@(private = "file")
ir_print_stmts :: proc(w: ^Writer, stmts: []IR_Stmt) {
	for stmt in stmts {
		ir_print_stmt(w, stmt)
	}
}

@(private = "file")
ir_print_stmt :: proc(w: ^Writer, stmt: IR_Stmt) {
	switch s in stmt {
	case ^IR_Let:
		for _ in 0 ..< w.indent do write(w, "\t")
		write(w, "let ", s.name, ": ", type_to_string(s.type), " = ")
		ir_print_expr(w, s.value)
		write(w, "\n")

	case ^IR_Assign:
		for _ in 0 ..< w.indent do write(w, "\t")
		ir_print_expr(w, s.target)
		write(w, " = ")
		ir_print_expr(w, s.value)
		write(w, "\n")

	case ^IR_Return:
		for _ in 0 ..< w.indent do write(w, "\t")
		write(w, "return")
		if s.value != nil {
			write(w, " ")
			ir_print_expr(w, s.value)
		}
		write(w, "\n")

	case ^IR_If:
		for _ in 0 ..< w.indent do write(w, "\t")
		write(w, "if ")
		ir_print_expr(w, s.condition)
		write(w, " then\n")
		indent(w)
		ir_print_stmts(w, s.then_body[:])
		dedent(w)
		for ei in s.elseif_clauses {
			for _ in 0 ..< w.indent do write(w, "\t")
			write(w, "elseif ")
			ir_print_expr(w, ei.condition)
			write(w, " then\n")
			indent(w)
			ir_print_stmts(w, ei.body[:])
			dedent(w)
		}
		if len(s.else_body) > 0 {
			write_line(w, "else")
			indent(w)
			ir_print_stmts(w, s.else_body[:])
			dedent(w)
		}
		write_line(w, "end")

	case ^IR_For:
		for _ in 0 ..< w.indent do write(w, "\t")
		write(w, "for ", s.var_name, " = ")
		ir_print_expr(w, s.start)
		write(w, ", ")
		ir_print_expr(w, s.stop)
		if s.step != nil {
			write(w, ", ")
			ir_print_expr(w, s.step)
		}
		write(w, " do\n")
		indent(w)
		ir_print_stmts(w, s.body[:])
		dedent(w)
		write_line(w, "end")

	case ^IR_While:
		for _ in 0 ..< w.indent do write(w, "\t")
		write(w, "while ")
		ir_print_expr(w, s.condition)
		write(w, " do\n")
		indent(w)
		ir_print_stmts(w, s.body[:])
		dedent(w)
		write_line(w, "end")

	case ^IR_Store_Output:
		for _ in 0 ..< w.indent do write(w, "\t")
		write(w, "store_output[", s.io_index, "] = ")
		ir_print_expr(w, s.value)
		write(w, "\n")

	case ^IR_Expr_Stmt:
		for _ in 0 ..< w.indent do write(w, "\t")
		ir_print_expr(w, s.expr)
		write(w, "\n")

	case ^IR_Barrier:
		for _ in 0 ..< w.indent do write(w, "\t")
		write_line(w, "barrier()")

	case ^IR_Discard:
		for _ in 0 ..< w.indent do write(w, "\t")
		write_line(w, "discard")

	case ^IR_Break:
		for _ in 0 ..< w.indent do write(w, "\t")
		write_line(w, "break")

	case ^IR_Continue:
		for _ in 0 ..< w.indent do write(w, "\t")
		write_line(w, "continue")
	}
}

@(private = "file")
ir_print_expr :: proc(w: ^Writer, expr: ^IR_Expr) {
	if expr == nil {
		write(w, "<nil>")
		return
	}

	switch e in expr.derived {
	case ^IR_Literal:
		switch v in e.value {
		case i64:  write(w, v)
		case f64:  write(w, v)
		case bool: write(w, v ? "true" : "false")
		}

	case ^IR_Var_Ref:
		write(w, e.name)

	case ^IR_Binary:
		write(w, "(")
		ir_print_expr(w, e.left)
		write(w, " ", ir_op_to_string(e.op), " ")
		ir_print_expr(w, e.right)
		write(w, ")")

	case ^IR_Unary:
		if e.op == .Neg {
			write(w, "(-")
		} else {
			write(w, "(!")
		}
		ir_print_expr(w, e.operand)
		write(w, ")")

	case ^IR_Call:
		if e.is_builtin do write(w, "@")
		write(w, e.name, "(")
		for arg, i in e.args {
			if i > 0 do write(w, ", ")
			ir_print_expr(w, arg)
		}
		write(w, ")")

	case ^IR_Field_Access:
		ir_print_expr(w, e.object)
		write(w, ".", e.field_name)

	case ^IR_Swizzle:
		ir_print_expr(w, e.object)
		write(w, ".", e.components)

	case ^IR_Index:
		ir_print_expr(w, e.object)
		write(w, "[")
		ir_print_expr(w, e.index)
		write(w, "]")

	case ^IR_Construct:
		write(w, e.type_name, "(")
		for arg, i in e.args {
			if i > 0 do write(w, ", ")
			ir_print_expr(w, arg)
		}
		write(w, ")")

	case ^IR_Type_Cast:
		write(w, type_to_string(expr.type), "(")
		ir_print_expr(w, e.value)
		write(w, ")")

	case ^IR_Load_Binding:
		write(w, "binding:", e.name)

	case ^IR_Input_Field:
		write(w, "input:", e.param_name, ".", e.field_name)

	case ^IR_Builtin_Var:
		write(w, "builtin:", e.name)

	case ^IR_Composite_Extract:
		ir_print_expr(w, e.object)
		write(w, ".", e.field_name)

	case ^IR_Vector_Shuffle:
		ir_print_expr(w, e.object)
		write(w, ".", swizzle_indices_to_string(e.components))

	case ^IR_Shared_Ref:
		write(w, "shared:", e.name)

	case ^IR_Select:
		write(w, "select(")
		ir_print_expr(w, e.condition)
		write(w, ", ")
		ir_print_expr(w, e.true_val)
		write(w, ", ")
		ir_print_expr(w, e.false_val)
		write(w, ")")
	}
}

ir_op_to_string :: proc(op: IR_Op) -> string {
	switch op {
	case .Add: return "+"
	case .Sub: return "-"
	case .Mul: return "*"
	case .Div: return "/"
	case .Mod: return "%"
	case .Eq:  return "=="
	case .Neq: return "!="
	case .Lt:  return "<"
	case .Gt:  return ">"
	case .Lte: return "<="
	case .Gte: return ">="
	case .And: return "&&"
	case .Or:  return "||"
	case .Neg: return "-"
	case .Not: return "!"
	}
	return "?"
}