package shader

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

// Shader reflection API — extract binding layouts, I/O, and struct info from IR

Reflect_Info :: struct {
	entry_points:   [dynamic]Reflect_Entry,
	structs:        [dynamic]Reflect_Struct,
	bindings:       [dynamic]Reflect_Binding,
	spec_constants: [dynamic]Reflect_Spec_Constant,
}

Reflect_Spec_Constant :: struct {
	name:          string,
	spec_id:       int,
	type_str:      string,
	default_value: string,
}

Reflect_Entry :: struct {
	name:       string,
	stage:      Shader_Stage,
	inputs:     [dynamic]Reflect_IO,
	outputs:    [dynamic]Reflect_IO,
	workgroup:  [3]int,
}

Reflect_IO :: struct {
	name:     string,
	type_str: string,
	location: int,
	builtin:  string,
}

Reflect_Binding :: struct {
	name:        string,
	group:       int,
	binding:     int,
	kind:        Reflect_Binding_Type,
	struct_name: string,
	size:        int,
	combined_id: int, // -1 if standalone, shared ID for split texture+sampler pairs
}

Reflect_Binding_Type :: enum {
	Uniform_Buffer,
	Storage_Buffer,
	Texture,
	Sampler,
	Push_Constant,
}

Reflect_Struct :: struct {
	name:   string,
	fields: [dynamic]Reflect_Struct_Field,
	size:   int,
}

Reflect_Struct_Field :: struct {
	name:     string,
	type_str: string,
	offset:   int,
	size:     int,
}

// Extract reflection info from an IR module
ir_reflect :: proc(module: ^IR_Module, allocator := context.allocator) -> Reflect_Info {
	info := Reflect_Info{
		entry_points   = make([dynamic]Reflect_Entry, allocator),
		structs        = make([dynamic]Reflect_Struct, allocator),
		bindings       = make([dynamic]Reflect_Binding, allocator),
		spec_constants = make([dynamic]Reflect_Spec_Constant, allocator),
	}

	// Entry points
	for &fn in module.functions {
		if !fn.is_entry do continue

		entry := Reflect_Entry{
			name       = fn.name,
			stage      = fn.stage,
			inputs     = make([dynamic]Reflect_IO, allocator),
			outputs    = make([dynamic]Reflect_IO, allocator),
			workgroup  = fn.workgroup_size,
		}

		for io in fn.inputs {
			append(&entry.inputs, Reflect_IO{
				name     = io.name,
				type_str = type_to_string(io.type),
				location = io.location,
				builtin  = io.builtin,
			})
		}

		for io in fn.outputs {
			append(&entry.outputs, Reflect_IO{
				name     = io.name,
				type_str = type_to_string(io.type),
				location = io.location,
				builtin  = io.builtin,
			})
		}

		append(&info.entry_points, entry)
	}

	// Bindings
	combined_counter := 0
	combined_ids: map[string]int // combined_name -> id
	for &b in module.bindings {
		kind: Reflect_Binding_Type
		switch b.kind {
		case .Uniform:       kind = .Uniform_Buffer
		case .Buffer:        kind = .Storage_Buffer
		case .Texture:       kind = .Texture
		case .Sampler:       kind = .Sampler
		case .Push_Constant: kind = .Push_Constant
		}

		struct_name := ""
		size := 0
		if b.struct_ref != nil {
			struct_name = b.struct_ref.name
			size = type_size_std140(b.type)
		}

		cid := -1
		if b.combined_name != "" {
			if existing, ok := combined_ids[b.combined_name]; ok {
				cid = existing
			} else {
				cid = combined_counter
				combined_ids[b.combined_name] = cid
				combined_counter += 1
			}
		}

		append(&info.bindings, Reflect_Binding{
			name        = b.name,
			group       = b.group,
			binding     = b.binding_num,
			kind        = kind,
			struct_name = struct_name,
			size        = size,
			combined_id = cid,
		})
	}

	// Structs
	for &s in module.structs {
		rs := Reflect_Struct{
			name   = s.name,
			fields = make([dynamic]Reflect_Struct_Field, allocator),
		}

		offset := 0
		for f in s.fields {
			field_size := type_size_std140(f.type)
			field_align := type_align_std140(f.type)

			// Align offset
			if field_align > 0 {
				offset = (offset + field_align - 1) & ~int(field_align - 1)
			}

			append(&rs.fields, Reflect_Struct_Field{
				name     = f.name,
				type_str = type_to_string(f.type),
				offset   = offset,
				size     = field_size,
			})

			offset += field_size
		}

		// Round up to alignment of the struct (vec4 = 16 bytes)
		rs.size = (offset + 15) & ~int(15)

		append(&info.structs, rs)
	}

	// Spec constants
	for &sc in module.spec_constants {
		append(&info.spec_constants, Reflect_Spec_Constant{
			name          = sc.name,
			spec_id       = sc.spec_id,
			type_str      = type_to_string(sc.type),
			default_value = ir_const_value_to_string(sc.default_value),
		})
	}

	return info
}

// Serialize reflection info to JSON
reflect_to_json :: proc(info: Reflect_Info, allocator := context.allocator) -> string {
	w := writer_init(allocator)

	write_line(&w, "{")
	indent(&w)

	// Entry points
	write_line(&w, "\"entry_points\": [")
	indent(&w)
	for entry, i in info.entry_points {
		write_line(&w, "{")
		indent(&w)

		stage_str: string
		#partial switch entry.stage {
		case .Vertex:   stage_str = "vertex"
		case .Fragment: stage_str = "fragment"
		case .Compute:  stage_str = "compute"
		case:           stage_str = "unknown"
		}

		write_line(&w, "\"name\": \"", entry.name, "\",")
		write_line(&w, "\"stage\": \"", stage_str, "\",")

		// Inputs
		write_line(&w, "\"inputs\": [")
		indent(&w)
		for io, j in entry.inputs {
			for _ in 0 ..< w.indent do write(&w, "\t")
			write(&w, "{ \"name\": \"", io.name, "\", \"type\": \"", io.type_str, "\"")
			if io.builtin != "" {
				write(&w, ", \"builtin\": \"", io.builtin, "\"")
			} else {
				write(&w, ", \"location\": ", io.location)
			}
			write(&w, " }")
			if j < len(entry.inputs) - 1 do write(&w, ",")
			write(&w, "\n")
		}
		dedent(&w)
		write_line(&w, "],")

		// Outputs
		write_line(&w, "\"outputs\": [")
		indent(&w)
		for io, j in entry.outputs {
			for _ in 0 ..< w.indent do write(&w, "\t")
			write(&w, "{ \"name\": \"", io.name, "\", \"type\": \"", io.type_str, "\"")
			if io.builtin != "" {
				write(&w, ", \"builtin\": \"", io.builtin, "\"")
			} else {
				write(&w, ", \"location\": ", io.location)
			}
			write(&w, " }")
			if j < len(entry.outputs) - 1 do write(&w, ",")
			write(&w, "\n")
		}
		dedent(&w)

		if entry.stage == .Compute {
			write_line(&w, "],")
			write_line(&w, "\"workgroup_size\": [", entry.workgroup[0], ", ", entry.workgroup[1], ", ", entry.workgroup[2], "]")
		} else {
			write_line(&w, "]")
		}

		dedent(&w)
		if i < len(info.entry_points) - 1 {
			write_line(&w, "},")
		} else {
			write_line(&w, "}")
		}
	}
	dedent(&w)
	write_line(&w, "],")

	// Bindings
	write_line(&w, "\"bindings\": [")
	indent(&w)
	for b, i in info.bindings {
		kind_str: string
		switch b.kind {
		case .Uniform_Buffer:  kind_str = "uniform_buffer"
		case .Storage_Buffer:  kind_str = "storage_buffer"
		case .Texture:         kind_str = "texture"
		case .Sampler:         kind_str = "sampler"
		case .Push_Constant:   kind_str = "push_constant"
		}

		for _ in 0 ..< w.indent do write(&w, "\t")
		write(&w, "{ \"name\": \"", b.name, "\", \"group\": ", b.group, ", \"binding\": ", b.binding, ", \"kind\": \"", kind_str, "\"")
		if b.struct_name != "" {
			write(&w, ", \"struct\": \"", b.struct_name, "\", \"size\": ", b.size)
		}
		if b.combined_id >= 0 {
			write(&w, ", \"combined_id\": ", b.combined_id)
		}
		write(&w, " }")
		if i < len(info.bindings) - 1 do write(&w, ",")
		write(&w, "\n")
	}
	dedent(&w)
	write_line(&w, "],")

	// Structs
	write_line(&w, "\"structs\": [")
	indent(&w)
	for s, i in info.structs {
		write_line(&w, "{")
		indent(&w)
		write_line(&w, "\"name\": \"", s.name, "\",")
		write_line(&w, "\"size\": ", s.size, ",")
		write_line(&w, "\"fields\": [")
		indent(&w)
		for f, j in s.fields {
			for _ in 0 ..< w.indent do write(&w, "\t")
			write(&w, "{ \"name\": \"", f.name, "\", \"type\": \"", f.type_str, "\", \"offset\": ", f.offset, ", \"size\": ", f.size, " }")
			if j < len(s.fields) - 1 do write(&w, ",")
			write(&w, "\n")
		}
		dedent(&w)
		write_line(&w, "]")
		dedent(&w)
		if i < len(info.structs) - 1 {
			write_line(&w, "},")
		} else {
			write_line(&w, "}")
		}
	}
	dedent(&w)
	if len(info.spec_constants) > 0 {
		write_line(&w, "],")
	} else {
		write_line(&w, "]")
	}

	// Spec constants
	if len(info.spec_constants) > 0 {
		write_line(&w, "\"spec_constants\": [")
		indent(&w)
		for sc, i in info.spec_constants {
			for _ in 0 ..< w.indent do write(&w, "\t")
			write(&w, "{ \"name\": \"", sc.name, "\", \"spec_id\": ", sc.spec_id, ", \"type\": \"", sc.type_str, "\", \"default\": ", sc.default_value, " }")
			if i < len(info.spec_constants) - 1 do write(&w, ",")
			write(&w, "\n")
		}
		dedent(&w)
		write_line(&w, "]")
	}

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

	return writer_to_string(w)
}

// std140 layout helpers

type_size_std140 :: proc(t: ^Resolved_Type) -> int {
	if t == nil do return 0
	switch v in t^ {
	case Type_Scalar:
		return 4
	case Type_Vector:
		switch v.size {
		case 2: return 8
		case 3: return 12
		case 4: return 16
		}
	case Type_Matrix:
		// Each column is a vec4 (std140 pads vec3 columns to vec4)
		return v.cols * 16
	case Type_Struct_Resolved:
		total := 0
		for f in v.fields {
			sz := type_size_std140(f.type)
			al := type_align_std140(f.type)
			if al > 0 {
				total = (total + al - 1) & ~(al - 1)
			}
			total += sz
		}
		return (total + 15) & ~int(15)
	case Type_Array_Resolved:
		elem_sz := type_size_std140(v.elem)
		// std140: array stride is rounded up to vec4
		stride := (elem_sz + 15) & ~int(15)
		return stride * v.size
	case Type_Sampler:
		return 0 // samplers don't have a buffer size
	case Type_Void:
		return 0
	}
	return 0
}

type_align_std140 :: proc(t: ^Resolved_Type) -> int {
	if t == nil do return 0
	switch v in t^ {
	case Type_Scalar:
		return 4
	case Type_Vector:
		switch v.size {
		case 2: return 8
		case 3: return 16 // vec3 aligned to vec4 in std140
		case 4: return 16
		}
	case Type_Matrix:
		return 16 // columns aligned to vec4
	case Type_Struct_Resolved:
		return 16 // structs aligned to vec4 in std140
	case Type_Array_Resolved:
		return 16 // arrays aligned to vec4 in std140
	case Type_Sampler:
		return 0
	case Type_Void:
		return 0
	}
	return 0
}

// -- SPIR-V-specific reflection --

Reflect_SPIRV_Info :: struct {
	base:            Reflect_Info,
	capabilities:    [dynamic]string,
	memory_model:    string,
	descriptor_sets: [dynamic]Reflect_Descriptor_Set,
}

Reflect_Descriptor_Set :: struct {
	set:      int,
	bindings: [dynamic]Reflect_SPIRV_Binding,
}

Reflect_SPIRV_Binding :: struct {
	binding:         int,
	name:            string,
	descriptor_type: string,
	stage_flags:     [dynamic]string,
	block_size:      int,
	members:         [dynamic]Reflect_SPIRV_Member,
}

Reflect_SPIRV_Member :: struct {
	name:          string,
	type_str:      string,
	offset:        int,
	size:          int,
	array_stride:  int,
	matrix_stride: int,
}

ir_reflect_spirv :: proc(module: ^IR_Module, allocator := context.allocator) -> Reflect_SPIRV_Info {
	info := Reflect_SPIRV_Info{
		base            = ir_reflect(module, allocator),
		capabilities    = make([dynamic]string, allocator),
		memory_model    = "GLSL450",
		descriptor_sets = make([dynamic]Reflect_Descriptor_Set, allocator),
	}

	// Capabilities — always Shader for now
	append(&info.capabilities, "Shader")

	// Build stage flags per binding: which entry points reference each binding
	binding_stages: map[string][dynamic]string
	for &fn in module.functions {
		if !fn.is_entry do continue
		stage_str: string
		#partial switch fn.stage {
		case .Vertex:   stage_str = "VERTEX"
		case .Fragment: stage_str = "FRAGMENT"
		case .Compute:  stage_str = "COMPUTE"
		case:           stage_str = "ALL"
		}
		// All bindings are potentially referenced by all entry points
		// (conservative — a real impl would do usage analysis)
		for &b in module.bindings {
			if binding_stages[b.name] == nil {
				binding_stages[b.name] = make([dynamic]string, allocator)
			}
			stages := &binding_stages[b.name]
			// Deduplicate
			found := false
			for s in stages {
				if s == stage_str { found = true; break }
			}
			if !found do append(stages, stage_str)
		}
	}

	// Group bindings by descriptor set
	set_map: map[int]int // group -> index in descriptor_sets
	for &b in module.bindings {
		set_idx: int
		if idx, ok := set_map[b.group]; ok {
			set_idx = idx
		} else {
			set_idx = len(info.descriptor_sets)
			set_map[b.group] = set_idx
			append(&info.descriptor_sets, Reflect_Descriptor_Set{
				set      = b.group,
				bindings = make([dynamic]Reflect_SPIRV_Binding, allocator),
			})
		}

		desc_type: string
		switch b.kind {
		case .Uniform:       desc_type = "UNIFORM_BUFFER"
		case .Buffer:        desc_type = "STORAGE_BUFFER"
		case .Texture:       desc_type = "SAMPLED_IMAGE"
		case .Sampler:       desc_type = "SAMPLER"
		case .Push_Constant: desc_type = "PUSH_CONSTANT"
		}

		block_size := 0
		members := make([dynamic]Reflect_SPIRV_Member, allocator)
		if b.struct_ref != nil {
			block_size = type_size_std140(b.type)
			offset := 0
			for f in b.struct_ref.fields {
				field_size := type_size_std140(f.type)
				field_align := type_align_std140(f.type)
				if field_align > 0 {
					offset = (offset + field_align - 1) & ~int(field_align - 1)
				}

				arr_stride := 0
				if ar, ok := f.type^.(Type_Array_Resolved); ok {
					elem_sz := type_size_std140(ar.elem)
					arr_stride = (elem_sz + 15) & ~int(15)
				}

				mat_stride := 0
				if _, ok := f.type^.(Type_Matrix); ok {
					mat_stride = 16
				}

				append(&members, Reflect_SPIRV_Member{
					name          = f.name,
					type_str      = type_to_string(f.type),
					offset        = offset,
					size          = field_size,
					array_stride  = arr_stride,
					matrix_stride = mat_stride,
				})
				offset += field_size
			}
		}

		stages := binding_stages[b.name] if b.name in binding_stages else make([dynamic]string, allocator)

		append(&info.descriptor_sets[set_idx].bindings, Reflect_SPIRV_Binding{
			binding         = b.binding_num,
			name            = b.name,
			descriptor_type = desc_type,
			stage_flags     = stages,
			block_size      = block_size,
			members         = members,
		})
	}

	return info
}

reflect_spirv_to_json :: proc(info: Reflect_SPIRV_Info, allocator := context.allocator) -> string {
	w := writer_init(allocator)

	write_line(&w, "{")
	indent(&w)

	// Base reflection (inline the entry_points, bindings, structs, spec_constants)
	base_json := reflect_to_json(info.base, allocator)
	// Strip outer { } and embed contents
	trimmed := strings.trim_space(base_json)
	if len(trimmed) > 2 {
		inner := trimmed[1:len(trimmed) - 1] // strip { and }
		inner = strings.trim_space(inner)
		write(&w, inner)
		write(&w, ",\n")
	}

	// Capabilities
	write_line(&w, "\"capabilities\": [")
	indent(&w)
	for cap, i in info.capabilities {
		for _ in 0 ..< w.indent do write(&w, "\t")
		write(&w, "\"", cap, "\"")
		if i < len(info.capabilities) - 1 do write(&w, ",")
		write(&w, "\n")
	}
	dedent(&w)
	write_line(&w, "],")

	// Memory model
	write_line(&w, "\"memory_model\": \"", info.memory_model, "\",")

	// Descriptor sets
	write_line(&w, "\"descriptor_sets\": [")
	indent(&w)
	for ds, di in info.descriptor_sets {
		write_line(&w, "{")
		indent(&w)
		write_line(&w, "\"set\": ", ds.set, ",")
		write_line(&w, "\"bindings\": [")
		indent(&w)
		for sb, bi in ds.bindings {
			write_line(&w, "{")
			indent(&w)
			write_line(&w, "\"binding\": ", sb.binding, ",")
			write_line(&w, "\"name\": \"", sb.name, "\",")
			write_line(&w, "\"descriptor_type\": \"", sb.descriptor_type, "\",")

			// Stage flags
			write(&w, "")
			for _ in 0 ..< w.indent do write(&w, "\t")
			write(&w, "\"stage_flags\": [")
			for sf, si in sb.stage_flags {
				write(&w, "\"", sf, "\"")
				if si < len(sb.stage_flags) - 1 do write(&w, ", ")
			}
			write(&w, "]")

			if sb.block_size > 0 {
				write(&w, ",\n")
				write_line(&w, "\"block_size\": ", sb.block_size, ",")
				write_line(&w, "\"members\": [")
				indent(&w)
				for m, mi in sb.members {
					for _ in 0 ..< w.indent do write(&w, "\t")
					write(&w, "{ \"name\": \"", m.name, "\", \"type\": \"", m.type_str, "\", \"offset\": ", m.offset, ", \"size\": ", m.size)
					if m.array_stride > 0 {
						write(&w, ", \"array_stride\": ", m.array_stride)
					}
					if m.matrix_stride > 0 {
						write(&w, ", \"matrix_stride\": ", m.matrix_stride)
					}
					write(&w, " }")
					if mi < len(sb.members) - 1 do write(&w, ",")
					write(&w, "\n")
				}
				dedent(&w)
				write_line(&w, "]")
			} else {
				write(&w, "\n")
			}

			dedent(&w)
			if bi < len(ds.bindings) - 1 {
				write_line(&w, "},")
			} else {
				write_line(&w, "}")
			}
		}
		dedent(&w)
		write_line(&w, "]")
		dedent(&w)
		if di < len(info.descriptor_sets) - 1 {
			write_line(&w, "},")
		} else {
			write_line(&w, "}")
		}
	}
	dedent(&w)
	write_line(&w, "]")

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

	return writer_to_string(w)
}