package shader

import "core:fmt"

// Desugar pass — runs after parsing, before semantic analysis.
// Transforms inline entry point syntax sugar into explicit struct form.

desugar_module :: proc(mod: ^Ast_Module) {
	for fn in mod.functions {
		if !is_entry_function(fn) do continue
		desugar_inline_entry(mod, fn)
	}
}

@(private = "file")
is_entry_function :: proc(fn: ^Ast_Function) -> bool {
	for attr in fn.attributes {
		if attr.name == "entry" do return true
	}
	return false
}

@(private = "file")
desugar_inline_entry :: proc(mod: ^Ast_Module, fn: ^Ast_Function) {
	// Desugar return type -> anonymous output struct
	// Handles both tuple return (name: Type, ...) and plain type shorthand (-> vec4)
	out_struct_name: string
	out_struct_fields: []Ast_Struct_Field
	if fn.return_type != nil {
		if tt, ok := fn.return_type^.(Type_Tuple); ok {
			// Tuple return: -> (color: vec4, ...)
			out_struct_name = fmt.aprintf("_%s_Output", fn.name)
			out_struct_fields = tt.fields
			out_struct := new(Ast_Struct)
			out_struct^ = Ast_Struct{
				name   = out_struct_name,
				fields = tt.fields,
				span   = tt.span,
			}
			append(&mod.structs, out_struct)
			fn.return_type^ = Type_Named{name = out_struct_name, span = tt.span}
		} else if tn, ok := fn.return_type^.(Type_Named); ok && is_primitive_type_name(tn.name) {
			// Plain type shorthand: -> vec4  (desugars to single-field output struct)
			out_struct_name = fmt.aprintf("_%s_Output", fn.name)
			field := Ast_Struct_Field{name = "_value", type = tn, span = tn.span}
			out_struct_fields = make([]Ast_Struct_Field, 1)
			out_struct_fields[0] = field
			out_struct := new(Ast_Struct)
			out_struct^ = Ast_Struct{
				name   = out_struct_name,
				fields = out_struct_fields,
				span   = tn.span,
			}
			append(&mod.structs, out_struct)
			fn.return_type^ = Type_Named{name = out_struct_name, span = tn.span}
		}
	}

	// Implicit single-field return: rewrite `return expr` -> `return Struct { field = expr }`
	if len(out_struct_fields) == 1 {
		rewrite_returns_single_field(fn.body[:], out_struct_name, out_struct_fields[0].name)
	}

	// Desugar inline params -> anonymous input struct
	// An inline entry point has multiple params OR params with attributes
	// (vs the normal single-struct-param convention)
	if should_desugar_params(fn) {
		struct_name := fmt.aprintf("_%s_Input", fn.name)
		fields := make([dynamic]Ast_Struct_Field)
		for p in fn.params {
			append(&fields, Ast_Struct_Field{
				name       = p.name,
				type       = p.type,
				attributes = p.attributes,
				span       = p.span,
			})
		}
		in_struct := new(Ast_Struct)
		in_struct^ = Ast_Struct{
			name   = struct_name,
			fields = fields[:],
			span   = fn.span,
		}
		append(&mod.structs, in_struct)

		// Collect original param names for body rewriting
		param_names: map[string]bool
		for p in fn.params {
			param_names[p.name] = true
		}

		// Rewrite function params to a single struct param
		new_params := make([]Ast_Param, 1)
		new_params[0] = Ast_Param{
			name = "input",
			type = Type_Named{name = struct_name, span = fn.span},
			span = fn.span,
		}
		fn.params = new_params

		// Rewrite body references: bare `uv` -> `input.uv`
		for &stmt in fn.body {
			rewrite_node(stmt, &param_names)
		}
	}
}

// Rewrite bare identifier references to desugared param names into field accesses
@(private = "file")
rewrite_node :: proc(node: ^Ast_Node, param_names: ^map[string]bool) {
	if node == nil do return

	#partial switch d in node.derived {
	case ^Ast_Let:
		rewrite_expr_in_place(&d.value, param_names)
	case ^Ast_Assign:
		rewrite_expr_in_place(&d.target, param_names)
		rewrite_expr_in_place(&d.value, param_names)
	case ^Ast_Output_Assign:
		rewrite_expr_in_place(&d.value, param_names)
	case ^Ast_Return:
		rewrite_expr_in_place(&d.value, param_names)
	case ^Ast_If:
		rewrite_expr_in_place(&d.condition, param_names)
		for &s in d.then_body do rewrite_node(s, param_names)
		for &clause in d.elseif_clauses {
			rewrite_expr_in_place(&clause.condition, param_names)
			for &s in clause.body do rewrite_node(s, param_names)
		}
		for &s in d.else_body do rewrite_node(s, param_names)
	case ^Ast_For:
		rewrite_expr_in_place(&d.start, param_names)
		rewrite_expr_in_place(&d.stop, param_names)
		if d.step != nil do rewrite_expr_in_place(&d.step, param_names)
		for &s in d.body do rewrite_node(s, param_names)
	case ^Ast_While:
		rewrite_expr_in_place(&d.condition, param_names)
		for &s in d.body do rewrite_node(s, param_names)
	case ^Ast_Call:
		for &arg in d.args do rewrite_expr_in_place(&arg, param_names)
	case ^Ast_Field_Access:
		rewrite_expr_in_place(&d.object, param_names)
	case ^Ast_Index:
		rewrite_expr_in_place(&d.object, param_names)
		rewrite_expr_in_place(&d.index, param_names)
	case ^Ast_Unary:
		rewrite_expr_in_place(&d.operand, param_names)
	case ^Ast_Binary:
		rewrite_expr_in_place(&d.left, param_names)
		rewrite_expr_in_place(&d.right, param_names)
	case ^Ast_Swizzle:
		rewrite_expr_in_place(&d.object, param_names)
	case ^Ast_Struct_Literal:
		for &f in d.fields {
			rewrite_expr_in_place(&f.value, param_names)
		}
	}
}

@(private = "file")
rewrite_expr_in_place :: proc(node_ptr: ^^Ast_Node, param_names: ^map[string]bool) {
	node := node_ptr^
	if node == nil do return

	// If this is an identifier that matches a desugared param, rewrite to input.field
	if ident, ok := node.derived.(^Ast_Ident); ok {
		if ident.name in param_names^ {
			// Create input.field_name
			input_ident := new(Ast_Ident)
			input_ident^ = Ast_Ident{name = "input", span = node.span}
			input_node := make_node(.Ident, input_ident, node.span)

			fa := new(Ast_Field_Access)
			fa^ = Ast_Field_Access{object = input_node, field = ident.name, span = node.span}
			node_ptr^ = make_node(.Field_Access, fa, node.span)
			return
		}
	}

	// Recurse into sub-expressions
	rewrite_node(node, param_names)
}

// Rewrite `return expr` -> `return StructName { field_name = expr }` for single-field output structs.
// Skips returns that are already struct literals for the output type.
@(private = "file")
rewrite_returns_single_field :: proc(stmts: []^Ast_Node, struct_name: string, field_name: string) {
	for stmt in stmts {
		if stmt == nil do continue
		#partial switch d in stmt.derived {
		case ^Ast_Return:
			if d.value == nil do continue
			// Skip if already a struct literal for our output type
			if sl, ok := d.value.derived.(^Ast_Struct_Literal); ok {
				if sl.type_name == struct_name do continue
			}
			// Wrap: return expr -> return StructName { field_name = expr }
			lit_fields := make([]Ast_Struct_Literal_Field, 1)
			lit_fields[0] = Ast_Struct_Literal_Field{
				name  = field_name,
				value = d.value,
				span  = d.value.span,
			}
			sl := new(Ast_Struct_Literal)
			sl^ = Ast_Struct_Literal{
				type_name = struct_name,
				fields    = lit_fields,
				span      = d.value.span,
			}
			d.value = make_node(.Struct_Literal, sl, d.value.span)
		case ^Ast_If:
			rewrite_returns_single_field(d.then_body[:], struct_name, field_name)
			for &clause in d.elseif_clauses {
				rewrite_returns_single_field(clause.body[:], struct_name, field_name)
			}
			rewrite_returns_single_field(d.else_body[:], struct_name, field_name)
		case ^Ast_For:
			rewrite_returns_single_field(d.body[:], struct_name, field_name)
		case ^Ast_While:
			rewrite_returns_single_field(d.body[:], struct_name, field_name)
		}
	}
}

@(private = "file")
should_desugar_params :: proc(fn: ^Ast_Function) -> bool {
	if len(fn.params) == 0 do return false

	if has_attribute(fn.attributes, "luma_workflow") do return true

	// Multiple params means inline entry point
	if len(fn.params) > 1 do return true

	// Single param with attributes means inline entry point
	if len(fn.params) == 1 && len(fn.params[0].attributes) > 0 do return true

	// Single param with a primitive type (not a struct) means inline entry point
	// Check if the type is a known primitive/vector/matrix type
	if len(fn.params) == 1 {
		name := type_expr_name(fn.params[0].type)
		if is_primitive_type_name(name) do return true
	}

	return false
}

@(private = "file")
is_primitive_type_name :: proc(name: string) -> bool {
	switch name {
	case "bool", "int", "uint", "float", "half",
	     "vec2", "vec3", "vec4",
	     "ivec2", "ivec3", "ivec4",
	     "uvec2", "uvec3", "uvec4",
	     "bvec2", "bvec3", "bvec4",
	     "mat2", "mat3", "mat4",
	     "mat2x3", "mat2x4", "mat3x2", "mat3x4", "mat4x2", "mat4x3":
		return true
	}
	return false
}