package shader

import "core:fmt"
import "core:strings"

// GLSL 450 backend — emits from IR_Module

Glsl_Flavor :: enum {
	Vulkan,
	OpenGL,
}

Glsl_Options :: struct {
	flavor: Glsl_Flavor,
}

GLSL_OPENGL_BINDINGS_PER_GROUP :: 16
GLSL_OPENGL_PUSH_CONSTANT_BINDING :: 15

Glsl_Emitter :: struct {
	w:           Writer,
	module:      ^IR_Module,
	current_fn:  ^IR_Function,
	diagnostics: [dynamic]Diagnostic,
	flavor:      Glsl_Flavor,
}

emit_glsl :: proc(module: ^IR_Module, allocator := context.allocator) -> (string, []Diagnostic) {
	return emit_glsl_with_options(module, {}, allocator)
}

emit_glsl_opengl :: proc(module: ^IR_Module, allocator := context.allocator) -> (string, []Diagnostic) {
	return emit_glsl_with_options(module, {flavor = .OpenGL}, allocator)
}

@(private = "file")
emit_glsl_with_options :: proc(module: ^IR_Module, options: Glsl_Options, allocator := context.allocator) -> (string, []Diagnostic) {
	e := Glsl_Emitter{
		w           = writer_init(allocator),
		module      = module,
		diagnostics = make([dynamic]Diagnostic, allocator),
		flavor      = options.flavor,
	}

	write_line(&e.w, "#version 450")
	write_line(&e.w, "")

	// Specialization constants
	for &sc in module.spec_constants {
		if e.flavor == .OpenGL {
			write_line(&e.w, "const ", resolved_type_to_glsl(sc.type), " ", sc.name, " = ", ir_const_value_to_string(sc.default_value), ";")
		} else {
			// layout(constant_id = N) const type name = default;
			write_line(&e.w, "layout(constant_id = ", sc.spec_id, ") const ", resolved_type_to_glsl(sc.type), " ", sc.name, " = ", ir_const_value_to_string(sc.default_value), ";")
		}
	}
	if len(module.spec_constants) > 0 {
		write_line(&e.w, "")
	}

	// Shared variables (workgroup memory)
	for sv in module.shared_vars {
		base, suffix := glsl_type_and_array_suffix(sv.type)
		write_line(&e.w, "shared ", base, " ", sv.name, suffix, ";")
	}
	if len(module.shared_vars) > 0 {
		write_line(&e.w, "")
	}

	emitted_struct := false
	for &s in module.structs {
		if glsl_struct_is_nested_field(module, &s) {
			emit_glsl_struct_def(&e, &s)
			emitted_struct = true
		}
	}
	if emitted_struct {
		write_line(&e.w, "")
	}

	// Bindings
	for &b in module.bindings {
		emit_glsl_binding(&e, &b)
	}
	if len(module.bindings) > 0 {
		write_line(&e.w, "")
	}

	// Functions
	for &fn in module.functions {
		emit_glsl_function(&e, &fn)
		write_line(&e.w, "")
	}

	return writer_to_string(e.w), e.diagnostics[:]
}

@(private = "file")
emit_glsl_struct_def :: proc(e: ^Glsl_Emitter, s: ^IR_Struct) {
	write_line(&e.w, "struct ", s.name, " {")
	indent(&e.w)
	for f in s.fields {
		base, suffix := glsl_type_and_array_suffix(f.type)
		write_line(&e.w, base, " ", f.name, suffix, ";")
	}
	dedent(&e.w)
	write_line(&e.w, "};")
}

@(private = "file")
glsl_struct_is_nested_field :: proc(module: ^IR_Module, candidate: ^IR_Struct) -> bool {
	for &s in module.structs {
		if s.name == candidate.name {
			continue
		}
		for f in s.fields {
			if glsl_type_references_struct(f.type, candidate.name) {
				return true
			}
		}
	}
	return false
}

@(private = "file")
glsl_type_references_struct :: proc(t: ^Resolved_Type, name: string) -> bool {
	if t == nil do return false
	switch v in t^ {
	case Type_Struct_Resolved:
		return v.name == name
	case Type_Array_Resolved:
		return glsl_type_references_struct(v.elem, name)
	case Type_Scalar, Type_Vector, Type_Matrix, Type_Sampler, Type_Void:
		return false
	}
	return false
}

// -- Bindings --

@(private = "file")
emit_glsl_binding :: proc(e: ^Glsl_Emitter, b: ^IR_Binding) {
	switch b.kind {
	case .Texture:
		if e.flavor == .OpenGL {
			binding := glsl_opengl_binding_point(b)
			write_line(&e.w, "layout(binding = ", binding, ") uniform ", resolved_type_to_glsl(b.type), " ", b.combined_name, ";")
		} else {
			// GLSL 450 uses combined samplers — emit as sampler2D with the original name
			write_line(&e.w, "layout(set = ", b.group, ", binding = ", b.binding_num, ") uniform ", resolved_type_to_glsl(b.type), " ", b.combined_name, ";")
		}
	case .Sampler:
		// Skip — already emitted by the Texture binding as a combined sampler
		return
	case .Uniform:
		if e.flavor == .OpenGL {
			binding := glsl_opengl_binding_point(b)
			if binding == GLSL_OPENGL_PUSH_CONSTANT_BINDING {
				append(&e.diagnostics, Diagnostic{level = .Error, message = fmt.aprintf("GLSL OpenGL: uniform binding for '%s' collides with reserved push-constant binding %d", b.name, GLSL_OPENGL_PUSH_CONSTANT_BINDING)})
			}
			write_line(&e.w, "layout(std140, binding = ", binding, ") uniform ", b.name, "_Block {")
		} else {
			write_line(&e.w, "layout(set = ", b.group, ", binding = ", b.binding_num, ") uniform ", b.name, "_Block {")
		}
		indent(&e.w)
		if b.struct_ref != nil {
			for f in b.struct_ref.fields {
				base, suffix := glsl_type_and_array_suffix(f.type)
				write_line(&e.w, base, " ", f.name, suffix, ";")
			}
		} else {
			// Bare uniform (e.g. uniform light_dir: vec3) — wrap as single field
			base, suffix := glsl_type_and_array_suffix(b.type)
			write_line(&e.w, base, " ", b.name, suffix, ";")
		}
		dedent(&e.w)
		if b.struct_ref != nil {
			write_line(&e.w, "} ", b.name, ";")
		} else {
			// No instance name — field is accessed directly by binding name
			write_line(&e.w, "};")
		}
	case .Buffer:
		if e.flavor == .OpenGL {
			binding := glsl_opengl_binding_point(b)
			write_line(&e.w, "layout(std430, binding = ", binding, ") buffer ", b.name, "_Block {")
		} else {
			write_line(&e.w, "layout(set = ", b.group, ", binding = ", b.binding_num, ") buffer ", b.name, "_Block {")
		}
		if b.struct_ref != nil {
			indent(&e.w)
			for f in b.struct_ref.fields {
				base, suffix := glsl_type_and_array_suffix(f.type)
				write_line(&e.w, base, " ", f.name, suffix, ";")
			}
			dedent(&e.w)
		}
		write_line(&e.w, "} ", b.name, ";")
	case .Push_Constant:
		if e.flavor == .OpenGL {
			write_line(&e.w, "layout(std140, binding = ", GLSL_OPENGL_PUSH_CONSTANT_BINDING, ") uniform ", b.name, "_Block {")
		} else {
			write_line(&e.w, "layout(push_constant) uniform ", b.name, "_Block {")
		}
		if b.struct_ref != nil {
			indent(&e.w)
			for f in b.struct_ref.fields {
				base, suffix := glsl_type_and_array_suffix(f.type)
				write_line(&e.w, base, " ", f.name, suffix, ";")
			}
			dedent(&e.w)
		}
		write_line(&e.w, "} ", b.name, ";")
	}
}

// -- Functions --

@(private = "file")
emit_glsl_function :: proc(e: ^Glsl_Emitter, fn: ^IR_Function) {
	e.current_fn = fn

	if fn.is_entry {
		emit_glsl_entry_point(e, fn)
	} else {
		emit_glsl_helper_function(e, fn)
	}

	e.current_fn = nil
}

@(private = "file")
emit_glsl_entry_point :: proc(e: ^Glsl_Emitter, fn: ^IR_Function) {
	// Compute shader workgroup size
	if fn.stage == .Compute {
		ws := fn.workgroup_size
		write(&e.w, "layout(local_size_x = ", ws[0])
		if ws[1] > 0 do write(&e.w, ", local_size_y = ", ws[1])
		if ws[2] > 0 do write(&e.w, ", local_size_z = ", ws[2])
		write(&e.w, ") in;\n")
		write_line(&e.w, "")
	}

	// Emit flattened input/output variables (skip for compute)
	if fn.stage != .Compute {
		for io in fn.inputs {
			if io.builtin != "" do continue
			write_line(&e.w, "layout(location = ", io.location, ") in ", resolved_type_to_glsl(io.type), " in_", io.name, ";")
		}
		for io in fn.outputs {
			if io.builtin != "" do continue
			write_line(&e.w, "layout(location = ", io.location, ") out ", resolved_type_to_glsl(io.type), " out_", io.name, ";")
		}
		write_line(&e.w, "")
	}

	write_line(&e.w, "void main() {")
	indent(&e.w)

	emit_glsl_stmts(e, fn.body[:])

	dedent(&e.w)
	write_line(&e.w, "}")
}

@(private = "file")
emit_glsl_helper_function :: proc(e: ^Glsl_Emitter, fn: ^IR_Function) {
	ret_str := resolved_type_to_glsl(fn.return_type)

	params_buf := strings.builder_make()
	for p, i in fn.params {
		if i > 0 do strings.write_string(&params_buf, ", ")
		strings.write_string(&params_buf, resolved_type_to_glsl(p.type))
		strings.write_string(&params_buf, " ")
		strings.write_string(&params_buf, p.name)
	}

	write_line(&e.w, ret_str, " ", fn.name, "(", strings.to_string(params_buf), ") {")
	indent(&e.w)

	emit_glsl_stmts(e, fn.body[:])

	dedent(&e.w)
	write_line(&e.w, "}")
}

// -- Statements --

@(private = "file")
emit_glsl_stmts :: proc(e: ^Glsl_Emitter, stmts: []IR_Stmt) {
	for stmt in stmts {
		emit_glsl_stmt(e, stmt)
	}
}

@(private = "file")
emit_glsl_stmt :: proc(e: ^Glsl_Emitter, stmt: IR_Stmt) {
	switch s in stmt {
	case ^IR_Let:
		type_str := resolved_type_to_glsl(s.type)
		for _ in 0 ..< e.w.indent do write(&e.w, "\t")
		write(&e.w, type_str, " ", s.name, " = ")
		emit_glsl_expr(e, s.value)
		write(&e.w, ";\n")

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

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

	case ^IR_Store_Output:
		fn := e.current_fn
		if s.io_index >= 0 && s.io_index < len(fn.outputs) {
			io := fn.outputs[s.io_index]
			out_name: string
			if io.builtin != "" {
				out_name = builtin_to_glsl_variable(io.builtin, fn.stage, false, e.flavor)
			} else {
				out_name = fmt.aprintf("out_%s", io.name)
			}
			for _ in 0 ..< e.w.indent do write(&e.w, "\t")
			write(&e.w, out_name, " = ")
			emit_glsl_expr(e, s.value)
			write(&e.w, ";\n")
		}

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

	case ^IR_For:
		for _ in 0 ..< e.w.indent do write(&e.w, "\t")
		write(&e.w, "for (int ", s.var_name, " = ")
		emit_glsl_expr(e, s.start)
		write(&e.w, "; ", s.var_name, " <= ")
		emit_glsl_expr(e, s.stop)
		write(&e.w, "; ", s.var_name)
		if s.step != nil {
			write(&e.w, " += ")
			emit_glsl_expr(e, s.step)
		} else {
			write(&e.w, "++")
		}
		write(&e.w, ") {\n")
		indent(&e.w)
		emit_glsl_stmts(e, s.body[:])
		dedent(&e.w)
		write_line(&e.w, "}")

	case ^IR_While:
		for _ in 0 ..< e.w.indent do write(&e.w, "\t")
		write(&e.w, "while (")
		emit_glsl_expr(e, s.condition)
		write(&e.w, ") {\n")
		indent(&e.w)
		emit_glsl_stmts(e, s.body[:])
		dedent(&e.w)
		write_line(&e.w, "}")

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

	case ^IR_Barrier:
		write_line(&e.w, "barrier();")

	case ^IR_Discard:
		write_line(&e.w, "discard;")

	case ^IR_Break:
		write_line(&e.w, "break;")

	case ^IR_Continue:
		write_line(&e.w, "continue;")
	}
}

// -- Expressions --

@(private = "file")
emit_glsl_expr :: proc(e: ^Glsl_Emitter, expr: ^IR_Expr) {
	if expr == nil {
		write(&e.w, "/* nil */")
		return
	}

	switch d in expr.derived {
	case ^IR_Literal:
		switch v in d.value {
		case i64:
			write(&e.w, v)
		case f64:
			s := fmt.aprintf("%v", v)
			if !strings.contains(s, ".") && !strings.contains(s, "e") {
				write(&e.w, s, ".0")
			} else {
				write(&e.w, s)
			}
		case bool:
			write(&e.w, v ? "true" : "false")
		}

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

	case ^IR_Binary:
		write(&e.w, "(")
		emit_glsl_expr(e, d.left)
		write(&e.w, " ", ir_op_to_glsl(d.op), " ")
		emit_glsl_expr(e, d.right)
		write(&e.w, ")")

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

	case ^IR_Call:
		glsl_name := d.is_builtin ? builtin_to_glsl(d.name) : d.name
		// Handle shadow sampling — GLSL needs vec3(uv, ref) packing
		if d.is_builtin && d.name == "sample_shadow" && len(d.args) >= 3 {
			write(&e.w, "texture(")
			emit_glsl_expr(e, d.args[0])
			write(&e.w, ", vec3(")
			emit_glsl_expr(e, d.args[1])
			write(&e.w, ", ")
			emit_glsl_expr(e, d.args[2])
			write(&e.w, "))")
			return
		}
		write(&e.w, glsl_name, "(")
		for arg, i in d.args {
			if i > 0 do write(&e.w, ", ")
			emit_glsl_expr(e, arg)
		}
		write(&e.w, ")")

	case ^IR_Field_Access:
		emit_glsl_expr(e, d.object)
		write(&e.w, ".", d.field_name)

	case ^IR_Swizzle:
		emit_glsl_expr(e, d.object)
		write(&e.w, ".", d.components)

	case ^IR_Composite_Extract:
		emit_glsl_expr(e, d.object)
		write(&e.w, ".", d.field_name)

	case ^IR_Vector_Shuffle:
		emit_glsl_expr(e, d.object)
		write(&e.w, ".", swizzle_indices_to_string(d.components))

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

	case ^IR_Construct:
		glsl_name := type_name_to_glsl(d.type_name)
		write(&e.w, glsl_name, "(")
		for arg, i in d.args {
			if i > 0 do write(&e.w, ", ")
			emit_glsl_expr(e, arg)
		}
		write(&e.w, ")")

	case ^IR_Type_Cast:
		write(&e.w, resolved_type_to_glsl(expr.type), "(")
		emit_glsl_expr(e, d.value)
		write(&e.w, ")")

	case ^IR_Load_Binding:
		write(&e.w, d.name)

	case ^IR_Input_Field:
		write(&e.w, "in_", d.field_name)

	case ^IR_Builtin_Var:
		fn := e.current_fn
		write(&e.w, builtin_to_glsl_variable(d.name, fn != nil ? fn.stage : .None, d.is_input, e.flavor))

	case ^IR_Shared_Ref:
		write(&e.w, d.name)

	case ^IR_Select:
		write(&e.w, "(")
		emit_glsl_expr(e, d.condition)
		write(&e.w, " ? ")
		emit_glsl_expr(e, d.true_val)
		write(&e.w, " : ")
		emit_glsl_expr(e, d.false_val)
		write(&e.w, ")")
	}
}

// -- Helpers --

@(private = "file")
ir_op_to_glsl :: 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 "?"
}

builtin_to_glsl_variable :: proc(name: string, stage: Shader_Stage, is_input: bool, flavor: Glsl_Flavor = .Vulkan) -> string {
	switch name {
	case "position":
		if stage == .Vertex && !is_input do return "gl_Position"
		if stage == .Fragment && is_input do return "gl_FragCoord"
		return "gl_Position"
	case "vertex_id":
		if flavor == .OpenGL do return "gl_VertexID"
		return "gl_VertexIndex"
	case "instance_id":
		if flavor == .OpenGL do return "gl_InstanceID"
		return "gl_InstanceIndex"
	case "frag_coord":
		return "gl_FragCoord"
	case "front_facing":
		return "gl_FrontFacing"
	case "local_invocation_id":
		return "gl_LocalInvocationID"
	case "local_invocation_index":
		return "gl_LocalInvocationIndex"
	case "global_invocation_id":
		return "gl_GlobalInvocationID"
	case "workgroup_id":
		return "gl_WorkGroupID"
	}
	return name
}

glsl_opengl_binding_point :: proc(binding: ^IR_Binding) -> int {
	if binding == nil do return 0
	return binding.group * GLSL_OPENGL_BINDINGS_PER_GROUP + binding.binding_num
}

builtin_to_glsl :: proc(name: string) -> string {
	switch name {
	case "sample":          return "texture"
	case "sample_shadow":   return "texture"
	case "sample_level":    return "textureLod"
	case "sample_grad":     return "textureGrad"
	case "sample_compare":  return "texture"
	case "texel_fetch":     return "texelFetch"
	case "texture_size":    return "textureSize"
	case "atan2":           return "atan"
	case "dfdx":            return "dFdx"
	case "dfdy":            return "dFdy"
	}
	return name
}

resolved_type_to_glsl :: proc(t: ^Resolved_Type) -> string {
	if t == nil do return "void"
	switch v in t^ {
	case Type_Scalar:
		switch v.kind {
		case .Bool:  return "bool"
		case .Int:   return "int"
		case .Uint:  return "uint"
		case .Float: return "float"
		case .Half:  return "float"
		}
	case Type_Vector:
		prefix: string
		switch v.elem {
		case .Float: prefix = "vec"
		case .Int:   prefix = "ivec"
		case .Uint:  prefix = "uvec"
		case .Bool:  prefix = "bvec"
		case .Half:  prefix = "vec"
		}
		return fmt.aprintf("%s%d", prefix, v.size)
	case Type_Matrix:
		if v.cols == v.rows {
			return fmt.aprintf("mat%d", v.cols)
		}
		return fmt.aprintf("mat%dx%d", v.cols, v.rows)
	case Type_Struct_Resolved:
		return v.name
	case Type_Array_Resolved:
		elem_str := resolved_type_to_glsl(v.elem)
		if v.size == 0 {
			return fmt.aprintf("%s[]", elem_str)
		}
		return fmt.aprintf("%s[%d]", elem_str, v.size)
	case Type_Sampler:
		switch v.kind {
		case .Sampler2D:       return "sampler2D"
		case .Sampler3D:       return "sampler3D"
		case .SamplerCube:     return "samplerCube"
		case .Sampler2DArray:  return "sampler2DArray"
		case .Sampler2DShadow: return "sampler2DShadow"
		}
	case Type_Void:
		return "void"
	}
	return "void"
}

type_name_to_glsl :: proc(name: string) -> string {
	switch name {
	case "half": return "float"
	}
	return name
}

// Split a type into base type string and array suffix for GLSL declaration syntax.
// GLSL uses `type name[size]` not `type[size] name`.
@(private = "file")
glsl_type_and_array_suffix :: proc(t: ^Resolved_Type) -> (base: string, suffix: string) {
	if t == nil do return "void", ""
	#partial switch v in t^ {
	case Type_Array_Resolved:
		elem_str := resolved_type_to_glsl(v.elem)
		if v.size == 0 {
			return elem_str, "[]"
		}
		return elem_str, fmt.aprintf("[%d]", v.size)
	}
	return resolved_type_to_glsl(t), ""
}