package shader

Symbol_Kind :: enum {
	Variable,
	Parameter,
	Function,
	Struct_Type,
	Builtin_Function,
	Binding,
	Constant,
	Shared,
}

Symbol :: struct {
	name:    string,
	type:    ^Resolved_Type,
	kind:    Symbol_Kind,
	span:    Source_Span,
	mutable: bool,
}

Scope :: struct {
	parent:  ^Scope,
	symbols: map[string]^Symbol,
}

scope_new :: proc(parent: ^Scope = nil, allocator := context.allocator) -> ^Scope {
	s := new(Scope, allocator)
	s.parent = parent
	s.symbols = make(map[string]^Symbol, allocator = allocator)
	return s
}

scope_define :: proc(s: ^Scope, sym: ^Symbol) -> bool {
	if sym.name in s.symbols {
		return false // already defined in this scope
	}
	s.symbols[sym.name] = sym
	return true
}

scope_lookup :: proc(s: ^Scope, name: string) -> ^Symbol {
	current := s
	for current != nil {
		if sym, ok := current.symbols[name]; ok {
			return sym
		}
		current = current.parent
	}
	return nil
}

scope_lookup_local :: proc(s: ^Scope, name: string) -> ^Symbol {
	if sym, ok := s.symbols[name]; ok {
		return sym
	}
	return nil
}

// Builtin function overload

Builtin_Overload :: struct {
	param_types:  []^Resolved_Type,
	return_type:  ^Resolved_Type,
}

Builtin_Def :: struct {
	name:      string,
	overloads: []Builtin_Overload,
}

// Create the global scope with all builtins pre-populated
create_global_scope :: proc(allocator := context.allocator) -> ^Scope {
	s := scope_new(nil, allocator)

	// Register builtin type names as symbols
	register_type_symbol(s, "bool",  TYPE_BOOL, allocator)
	register_type_symbol(s, "int",   TYPE_INT, allocator)
	register_type_symbol(s, "uint",  TYPE_UINT, allocator)
	register_type_symbol(s, "float", TYPE_FLOAT, allocator)
	register_type_symbol(s, "half",  TYPE_HALF, allocator)
	register_type_symbol(s, "vec2",  TYPE_VEC2, allocator)
	register_type_symbol(s, "vec3",  TYPE_VEC3, allocator)
	register_type_symbol(s, "vec4",  TYPE_VEC4, allocator)
	register_type_symbol(s, "ivec2", TYPE_IVEC2, allocator)
	register_type_symbol(s, "ivec3", TYPE_IVEC3, allocator)
	register_type_symbol(s, "ivec4", TYPE_IVEC4, allocator)
	register_type_symbol(s, "uvec2", TYPE_UVEC2, allocator)
	register_type_symbol(s, "uvec3", TYPE_UVEC3, allocator)
	register_type_symbol(s, "uvec4", TYPE_UVEC4, allocator)
	register_type_symbol(s, "bvec2", TYPE_BVEC2, allocator)
	register_type_symbol(s, "bvec3", TYPE_BVEC3, allocator)
	register_type_symbol(s, "bvec4", TYPE_BVEC4, allocator)
	register_type_symbol(s, "mat2",  TYPE_MAT2, allocator)
	register_type_symbol(s, "mat3",  TYPE_MAT3, allocator)
	register_type_symbol(s, "mat4",  TYPE_MAT4, allocator)
	register_type_symbol(s, "sampler2D", TYPE_SAMPLER2D, allocator)
	register_type_symbol(s, "sampler2DShadow", TYPE_SAMPLER2D_SHADOW, allocator)

	// Register matrix variant names
	register_type_symbol(s, "mat2x2", TYPE_MAT2, allocator)
	register_type_symbol(s, "mat3x3", TYPE_MAT3, allocator)
	register_type_symbol(s, "mat4x4", TYPE_MAT4, allocator)

	// Additional matrix sizes
	mat2x3 := make_type(Type_Matrix{.Float, 2, 3}, allocator)
	mat2x4 := make_type(Type_Matrix{.Float, 2, 4}, allocator)
	mat3x2 := make_type(Type_Matrix{.Float, 3, 2}, allocator)
	mat3x4 := make_type(Type_Matrix{.Float, 3, 4}, allocator)
	mat4x2 := make_type(Type_Matrix{.Float, 4, 2}, allocator)
	mat4x3 := make_type(Type_Matrix{.Float, 4, 3}, allocator)
	register_type_symbol(s, "mat2x3", mat2x3, allocator)
	register_type_symbol(s, "mat2x4", mat2x4, allocator)
	register_type_symbol(s, "mat3x2", mat3x2, allocator)
	register_type_symbol(s, "mat3x4", mat3x4, allocator)
	register_type_symbol(s, "mat4x2", mat4x2, allocator)
	register_type_symbol(s, "mat4x3", mat4x3, allocator)

	// Register builtin functions
	// Math builtins (float -> float)
	math_1f := []string{
		"abs", "sign", "floor", "ceil", "round", "fract", "sqrt", "inversesqrt",
		"exp", "exp2", "log", "log2", "sin", "cos", "tan", "asin", "acos", "atan",
	}
	for name in math_1f {
		register_builtin_fn(s, name, TYPE_FLOAT, allocator)
	}

	// Math builtins (float, float -> float)
	math_2f := []string{"mod", "min", "max", "pow", "atan2", "step"}
	for name in math_2f {
		register_builtin_fn_2(s, name, TYPE_FLOAT, TYPE_FLOAT, TYPE_FLOAT, allocator)
	}

	// clamp, mix, smoothstep (3 args)
	register_builtin_fn_3(s, "clamp", TYPE_FLOAT, TYPE_FLOAT, TYPE_FLOAT, TYPE_FLOAT, allocator)
	register_builtin_fn_3(s, "mix", TYPE_FLOAT, TYPE_FLOAT, TYPE_FLOAT, TYPE_FLOAT, allocator)
	register_builtin_fn_3(s, "smoothstep", TYPE_FLOAT, TYPE_FLOAT, TYPE_FLOAT, TYPE_FLOAT, allocator)

	// Vector builtins
	register_builtin_fn(s, "normalize", TYPE_VEC3, allocator) // simplified - works on any vec
	register_builtin_fn(s, "length", TYPE_FLOAT, allocator)
	register_builtin_fn_2(s, "dot", TYPE_VEC3, TYPE_VEC3, TYPE_FLOAT, allocator)
	register_builtin_fn_2(s, "cross", TYPE_VEC3, TYPE_VEC3, TYPE_VEC3, allocator)
	register_builtin_fn_2(s, "distance", TYPE_VEC3, TYPE_VEC3, TYPE_FLOAT, allocator)
	register_builtin_fn_2(s, "reflect", TYPE_VEC3, TYPE_VEC3, TYPE_VEC3, allocator)
	register_builtin_fn_3(s, "refract", TYPE_VEC3, TYPE_VEC3, TYPE_FLOAT, TYPE_VEC3, allocator)
	register_builtin_fn_3(s, "faceforward", TYPE_VEC3, TYPE_VEC3, TYPE_VEC3, TYPE_VEC3, allocator)

	// Matrix builtins
	register_builtin_fn(s, "transpose", TYPE_MAT4, allocator)
	register_builtin_fn(s, "inverse", TYPE_MAT4, allocator)
	register_builtin_fn(s, "determinant", TYPE_FLOAT, allocator)

	// Texture builtins
	register_builtin_fn_2(s, "sample", TYPE_SAMPLER2D, TYPE_VEC2, TYPE_VEC4, allocator)
	register_builtin_fn_3(s, "sample_level", TYPE_SAMPLER2D, TYPE_VEC2, TYPE_FLOAT, TYPE_VEC4, allocator)
	register_builtin_fn_3(s, "sample_shadow", TYPE_SAMPLER2D_SHADOW, TYPE_VEC2, TYPE_FLOAT, TYPE_FLOAT, allocator)

	// Derivative builtins (fragment only)
	register_builtin_fn(s, "dfdx", TYPE_FLOAT, allocator)
	register_builtin_fn(s, "dfdy", TYPE_FLOAT, allocator)
	register_builtin_fn(s, "fwidth", TYPE_FLOAT, allocator)

	// Compute builtins
	register_builtin_fn(s, "barrier", TYPE_VOID, allocator)

	// Conditional select
	register_builtin_fn(s, "select", TYPE_VOID, allocator) // return type resolved specially in sema

	return s
}

@(private = "file")
register_type_symbol :: proc(s: ^Scope, name: string, t: ^Resolved_Type, allocator := context.allocator) {
	sym := new(Symbol, allocator)
	sym^ = Symbol{
		name = name,
		type = t,
		kind = .Struct_Type, // Using Struct_Type for all type symbols
	}
	scope_define(s, sym)
}

@(private = "file")
register_builtin_fn :: proc(s: ^Scope, name: string, ret_type: ^Resolved_Type, allocator := context.allocator) {
	sym := new(Symbol, allocator)
	sym^ = Symbol{
		name = name,
		type = ret_type,
		kind = .Builtin_Function,
	}
	scope_define(s, sym)
}

@(private = "file")
register_builtin_fn_2 :: proc(s: ^Scope, name: string, p1, p2, ret: ^Resolved_Type, allocator := context.allocator) {
	sym := new(Symbol, allocator)
	sym^ = Symbol{
		name = name,
		type = ret,
		kind = .Builtin_Function,
	}
	scope_define(s, sym)
}

@(private = "file")
register_builtin_fn_3 :: proc(s: ^Scope, name: string, p1, p2, p3, ret: ^Resolved_Type, allocator := context.allocator) {
	sym := new(Symbol, allocator)
	sym^ = Symbol{
		name = name,
		type = ret,
		kind = .Builtin_Function,
	}
	scope_define(s, sym)
}