package shader

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

Sema :: struct {
	global_scope:  ^Scope,
	current_scope: ^Scope,
	diagnostics:   [dynamic]Diagnostic,
	structs:       map[string]^Type_Struct_Resolved,
	functions:     map[string]^Ast_Function,
	bindings:      map[string]^Ast_Binding,
	current_stage: Shader_Stage,
	error_count:   int,
	loop_depth:    int,
}

sema_init :: proc(allocator := context.allocator) -> Sema {
	init_builtin_types()
	global := create_global_scope(allocator)
	return Sema{
		global_scope  = global,
		current_scope = global,
		diagnostics   = make([dynamic]Diagnostic, allocator),
		structs       = make(map[string]^Type_Struct_Resolved, allocator = allocator),
		functions     = make(map[string]^Ast_Function, allocator = allocator),
		bindings      = make(map[string]^Ast_Binding, allocator = allocator),
	}
}

check_module :: proc(s: ^Sema, mod: ^Ast_Module) -> []Diagnostic {
	// Pass 1: collect all declarations
	collect_declarations(s, mod)

	// Pass 2: check function bodies
	for fn in mod.functions {
		check_function(s, fn)
	}

	return s.diagnostics[:]
}

// -- Pass 1: Declaration collection --

@(private = "file")
collect_declarations :: proc(s: ^Sema, mod: ^Ast_Module) {
	// Register structs
	for st in mod.structs {
		fields := make([dynamic]Resolved_Struct_Field)
		for f in st.fields {
			rt := resolve_type_expr(s, f.type)
			append(&fields, Resolved_Struct_Field{
				name       = f.name,
				type       = rt,
				attributes = f.attributes,
			})
		}
		resolved := new(Type_Struct_Resolved)
		resolved^ = Type_Struct_Resolved{
			name   = st.name,
			fields = fields[:],
		}
		s.structs[st.name] = resolved

		// Register as type symbol
		t := make_type(resolved^)
		sym := new(Symbol)
		sym^ = Symbol{
			name = st.name,
			type = t,
			kind = .Struct_Type,
			span = st.span,
		}
		scope_define(s.global_scope, sym)
	}

	// Register bindings
	push_constant_count := 0
	for b in mod.bindings {
		s.bindings[b.name] = b
		rt := resolve_type_expr(s, b.type_expr)
		sym := new(Symbol)
		sym^ = Symbol{
			name = b.name,
			type = rt,
			kind = .Binding,
			span = b.span,
		}
		scope_define(s.global_scope, sym)

		// Validate @push_constant
		if has_attribute(b.attributes, "push_constant") {
			push_constant_count += 1
			if b.kind != .Uniform {
				sema_error(s, b.span, "@push_constant can only be applied to 'uniform', not 'buffer'")
			}
			eg, eb := get_group_binding(b.attributes)
			if eg >= 0 || eb >= 0 {
				sema_error(s, b.span, "@push_constant cannot have @group or @binding attributes")
			}
			if push_constant_count > 1 {
				sema_error(s, b.span, "only one @push_constant binding is allowed per module")
			}
		}
	}

	// Register constants
	for c in mod.constants {
		ct: ^Resolved_Type
		if c.type != nil {
			ct = resolve_type_expr(s, c.type^)
		} else if c.value != nil {
			ct = check_expr(s, c.value)
		}
		if ct != nil {
			sym := new(Symbol)
			sym^ = Symbol{
				name = c.name,
				type = ct,
				kind = .Constant,
				span = c.span,
			}
			scope_define(s.global_scope, sym)
		}
	}

	// Register shared variables
	for sv in mod.shared_vars {
		rt := resolve_type_expr(s, sv.type_expr)
		sym := new(Symbol)
		sym^ = Symbol{
			name    = sv.name,
			type    = rt,
			kind    = .Shared,
			span    = sv.span,
			mutable = true,
		}
		scope_define(s.global_scope, sym)
	}

	// Register functions
	for fn in mod.functions {
		s.functions[fn.name] = fn
		ret_type := TYPE_VOID
		if fn.return_type != nil {
			ret_type = resolve_type_expr(s, fn.return_type^)
		}
		sym := new(Symbol)
		sym^ = Symbol{
			name = fn.name,
			type = ret_type,
			kind = .Function,
			span = fn.span,
		}
		scope_define(s.global_scope, sym)
	}
}

// -- Pass 2: Function body checking --

@(private = "file")
check_function :: proc(s: ^Sema, fn: ^Ast_Function) {
	// Determine shader stage from attributes
	s.current_stage = .None
	for attr in fn.attributes {
		if attr.name == "entry" && len(attr.args) > 0 {
			switch attr.args[0] {
			case "vertex":   s.current_stage = .Vertex
			case "fragment": s.current_stage = .Fragment
			case "compute":  s.current_stage = .Compute
			}
		}
	}

	// Create function scope
	fn_scope := scope_new(s.global_scope)
	old_scope := s.current_scope
	s.current_scope = fn_scope

	// Register parameters
	for param in fn.params {
		pt := resolve_type_expr(s, param.type)
		sym := new(Symbol)
		sym^ = Symbol{
			name = param.name,
			type = pt,
			kind = .Parameter,
			span = param.span,
		}
		scope_define(fn_scope, sym)
	}

	// Check body
	for stmt in fn.body {
		check_stmt(s, stmt)
	}

	s.current_scope = old_scope
}

@(private = "file")
check_stmt :: proc(s: ^Sema, node: ^Ast_Node) {
	if node == nil do return

	#partial switch d in node.derived {
	case ^Ast_Let:
		check_let(s, d, node)
	case ^Ast_Assign:
		check_assign(s, d, node)
	case ^Ast_Output_Assign:
		check_output_assign(s, d, node)
	case ^Ast_Return:
		if d.value != nil {
			check_expr(s, d.value)
		}
	case ^Ast_If:
		check_if(s, d)
	case ^Ast_For:
		check_for(s, d)
	case ^Ast_While:
		check_while(s, d)
	case ^Ast_Discard:
		if s.current_stage != .Fragment {
			sema_error(s, d.span, "'discard' is only allowed in fragment shaders")
		}
	case ^Ast_Break:
		if s.loop_depth == 0 {
			sema_error(s, d.span, "'break' is only allowed inside a loop")
		}
	case ^Ast_Continue:
		if s.loop_depth == 0 {
			sema_error(s, d.span, "'continue' is only allowed inside a loop")
		}
	case ^Ast_Call:
		check_expr(s, node)
	case ^Ast_Ident:
		check_expr(s, node)
	case ^Ast_Binary:
		check_expr(s, node)
	case ^Ast_Unary:
		check_expr(s, node)
	case ^Ast_Field_Access:
		check_expr(s, node)
	case:
		// Expression statement
		check_expr(s, node)
	}
}

@(private = "file")
check_output_assign :: proc(s: ^Sema, assign: ^Ast_Output_Assign, node: ^Ast_Node) {
	output_type := resolve_type_expr(s, assign.type)
	value_type := check_expr(s, assign.value)
	if output_type != nil && value_type != nil {
		if !type_equals(output_type, value_type) && !can_coerce(value_type, output_type) {
			sema_error(s, assign.span, fmt.aprintf(
				"cannot assign %s to out slot '%s' of type %s",
				type_to_string(value_type),
				assign.name,
				type_to_string(output_type),
			))
		}
	}
	node.resolved_type = output_type
}

@(private = "file")
check_let :: proc(s: ^Sema, let: ^Ast_Let, node: ^Ast_Node) {
	value_type := check_expr(s, let.value)

	declared_type: ^Resolved_Type
	if let.type_expr != nil {
		declared_type = resolve_type_expr(s, let.type_expr^)
		if !type_equals(declared_type, value_type) && value_type != nil {
			// Allow numeric literal coercion
			if !can_coerce(value_type, declared_type) {
				sema_error(s, let.span, fmt.aprintf(
					"type mismatch: declared %s but got %s",
					type_to_string(declared_type),
					type_to_string(value_type),
				))
			}
		}
	} else {
		declared_type = value_type
	}

	sym := new(Symbol)
	sym^ = Symbol{
		name    = let.name,
		type    = declared_type,
		kind    = .Variable,
		span    = let.span,
		mutable = false,
	}
	scope_define(s.current_scope, sym)
	node.resolved_type = declared_type
}

@(private = "file")
check_assign :: proc(s: ^Sema, assign: ^Ast_Assign, node: ^Ast_Node) {
	target_type := check_expr(s, assign.target)
	value_type := check_expr(s, assign.value)
	if target_type != nil && value_type != nil {
		if !type_equals(target_type, value_type) && !can_coerce(value_type, target_type) {
			sema_error(s, assign.span, fmt.aprintf(
				"cannot assign %s to %s",
				type_to_string(value_type),
				type_to_string(target_type),
			))
		}
	}
}

@(private = "file")
check_if :: proc(s: ^Sema, if_node: ^Ast_If) {
	check_expr(s, if_node.condition)
	for stmt in if_node.then_body {
		check_stmt(s, stmt)
	}
	for ei in if_node.elseif_clauses {
		check_expr(s, ei.condition)
		for stmt in ei.body {
			check_stmt(s, stmt)
		}
	}
	for stmt in if_node.else_body {
		check_stmt(s, stmt)
	}
}

@(private = "file")
check_for :: proc(s: ^Sema, for_node: ^Ast_For) {
	check_expr(s, for_node.start)
	check_expr(s, for_node.stop)
	if for_node.step != nil {
		check_expr(s, for_node.step)
	}

	loop_scope := scope_new(s.current_scope)
	old := s.current_scope
	s.current_scope = loop_scope

	sym := new(Symbol)
	sym^ = Symbol{
		name = for_node.var_name,
		type = TYPE_INT,
		kind = .Variable,
		span = for_node.span,
	}
	scope_define(loop_scope, sym)

	s.loop_depth += 1
	for stmt in for_node.body {
		check_stmt(s, stmt)
	}
	s.loop_depth -= 1

	s.current_scope = old
}

@(private = "file")
check_while :: proc(s: ^Sema, while_node: ^Ast_While) {
	check_expr(s, while_node.condition)
	s.loop_depth += 1
	for stmt in while_node.body {
		check_stmt(s, stmt)
	}
	s.loop_depth -= 1
}

// -- Expression type checking --

@(private = "file")
check_expr :: proc(s: ^Sema, node: ^Ast_Node) -> ^Resolved_Type {
	if node == nil do return nil

	result: ^Resolved_Type

	#partial switch d in node.derived {
	case ^Ast_Literal:
		result = check_literal(s, d)
	case ^Ast_Ident:
		result = check_ident(s, d)
	case ^Ast_Binary:
		result = check_binary(s, d)
	case ^Ast_Unary:
		result = check_unary(s, d)
	case ^Ast_Call:
		result = check_call(s, d)
	case ^Ast_Field_Access:
		result = check_field_access(s, d)
	case ^Ast_Index:
		result = check_index(s, d)
	case ^Ast_Struct_Literal:
		result = check_struct_literal(s, d)
	case:
		result = nil
	}

	node.resolved_type = result
	return result
}

@(private = "file")
check_literal :: proc(s: ^Sema, lit: ^Ast_Literal) -> ^Resolved_Type {
	switch v in lit.value {
	case i64:    return TYPE_INT
	case f64:    return TYPE_FLOAT
	case bool:   return TYPE_BOOL
	case string: return nil // strings not really used at runtime
	}
	return nil
}

@(private = "file")
check_ident :: proc(s: ^Sema, ident: ^Ast_Ident) -> ^Resolved_Type {
	sym := scope_lookup(s.current_scope, ident.name)
	if sym == nil {
		sema_error(s, ident.span, fmt.aprintf("undefined identifier '%s'", ident.name))
		suggest_similar(s, ident.name, ident.span)
		return nil
	}
	return sym.type
}

@(private = "file")
check_binary :: proc(s: ^Sema, bin: ^Ast_Binary) -> ^Resolved_Type {
	left_type := check_expr(s, bin.left)
	right_type := check_expr(s, bin.right)

	if left_type == nil || right_type == nil do return nil

	switch bin.op {
	case .And, .Or:
		return TYPE_BOOL
	case .Eq, .Neq, .Lt, .Gt, .Lte, .Gte:
		return TYPE_BOOL
	case .Add, .Sub, .Mul, .Div, .Mod:
		return resolve_arithmetic_type(s, left_type, right_type, bin.left.span)
	}
	return nil
}

@(private = "file")
check_unary :: proc(s: ^Sema, un: ^Ast_Unary) -> ^Resolved_Type {
	operand_type := check_expr(s, un.operand)
	switch un.op {
	case .Neg:
		return operand_type
	case .Not:
		return TYPE_BOOL
	}
	return nil
}

@(private = "file")
check_call :: proc(s: ^Sema, call: ^Ast_Call) -> ^Resolved_Type {
	// Check arguments
	for arg in call.args {
		check_expr(s, arg)
	}

	// Get callee name
	if call.callee.kind == .Ident {
		ident := call.callee.derived.(^Ast_Ident)
		sym := scope_lookup(s.current_scope, ident.name)
		if sym == nil {
			sema_error(s, call.span, fmt.aprintf("undefined function '%s'", ident.name))
			suggest_similar(s, ident.name, call.span)
			return nil
		}

		// Handle vector/type constructors
		if sym.kind == .Struct_Type {
			return resolve_constructor_call(s, ident.name, call)
		}

		// Handle builtin and user functions
		if sym.kind == .Builtin_Function || sym.kind == .Function {
			return resolve_function_return_type(s, sym, call)
		}
	}

	check_expr(s, call.callee)
	return nil
}

@(private = "file")
resolve_constructor_call :: proc(s: ^Sema, name: string, call: ^Ast_Call) -> ^Resolved_Type {
	sym := scope_lookup(s.current_scope, name)
	if sym == nil do return nil

	// Check if it's a struct type
	if resolved_struct, ok := s.structs[name]; ok {
		// This is handled by struct literal, but constructors also work
		return sym.type
	}

	// It's a builtin type constructor (vec2, vec3, etc.)
	return sym.type
}

@(private = "file")
resolve_function_return_type :: proc(s: ^Sema, sym: ^Symbol, call: ^Ast_Call) -> ^Resolved_Type {
	if sym.kind == .Builtin_Function {
		// For builtins that return the same type as their first arg
		name := sym.name
		passthrough_builtins := []string{
			"abs", "sign", "floor", "ceil", "round", "fract", "sqrt",
			"min", "max", "clamp", "mix", "normalize", "reflect",
			"sin", "cos", "tan", "asin", "acos", "atan",
			"transpose", "inverse",
		}
		for b in passthrough_builtins {
			if name == b && len(call.args) > 0 {
				arg_type := call.args[0].resolved_type
				if arg_type != nil {
					return arg_type
				}
			}
		}
		// select(cond, true_val, false_val) -> type of true_val
		if name == "select" {
			if len(call.args) != 3 {
				sema_error(s, call.span, "'select' requires exactly 3 arguments: select(condition, true_val, false_val)")
				return nil
			}
			cond_type := call.args[0].resolved_type
			if cond_type != nil && cond_type != TYPE_BOOL {
				sema_error(s, call.span, "'select' condition must be bool")
			}
			true_type := call.args[1].resolved_type
			false_type := call.args[2].resolved_type
			if true_type != nil && false_type != nil && true_type != false_type {
				sema_error(s, call.span, "'select' true and false values must have matching types")
			}
			return true_type
		}
		return sym.type
	}

	// User function — return type is sym.type
	return sym.type
}

@(private = "file")
check_field_access :: proc(s: ^Sema, fa: ^Ast_Field_Access) -> ^Resolved_Type {
	obj_type := check_expr(s, fa.object)
	if obj_type == nil do return nil

	// Check for swizzle on vector types
	if v, ok := obj_type^.(Type_Vector); ok {
		if is_valid_swizzle(fa.field, v.size) {
			swizzle_len := len(fa.field)
			if swizzle_len == 1 {
				return make_type(Type_Scalar{v.elem})
			}
			return make_type(Type_Vector{v.elem, swizzle_len})
		}
	}

	// Check for struct field access
	#partial switch t in obj_type^ {
	case Type_Struct_Resolved:
		for f in t.fields {
			if f.name == fa.field {
				return f.type
			}
		}
		sema_error(s, fa.span, fmt.aprintf("no field '%s' in struct '%s'", fa.field, t.name))
	case:
		// Could be swizzle on non-vector — error
		sema_error(s, fa.span, fmt.aprintf("cannot access field '%s' on type %s", fa.field, type_to_string(obj_type)))
	}

	return nil
}

@(private = "file")
check_index :: proc(s: ^Sema, idx: ^Ast_Index) -> ^Resolved_Type {
	obj_type := check_expr(s, idx.object)
	check_expr(s, idx.index)

	if obj_type == nil do return nil

	#partial switch t in obj_type^ {
	case Type_Array_Resolved:
		return t.elem
	case Type_Vector:
		return make_type(Type_Scalar{t.elem})
	case Type_Matrix:
		return make_type(Type_Vector{t.elem, t.rows})
	case:
		sema_error(s, idx.span, "type is not indexable")
	}
	return nil
}

@(private = "file")
check_struct_literal :: proc(s: ^Sema, sl: ^Ast_Struct_Literal) -> ^Resolved_Type {
	resolved, ok := s.structs[sl.type_name]
	if !ok {
		sema_error(s, sl.span, fmt.aprintf("undefined struct type '%s'", sl.type_name))
		return nil
	}

	// Check fields
	for f in sl.fields {
		check_expr(s, f.value)

		// Verify field exists
		found := false
		for rf in resolved.fields {
			if rf.name == f.name {
				found = true
				break
			}
		}
		if !found {
			sema_error(s, f.span, fmt.aprintf("struct '%s' has no field '%s'", sl.type_name, f.name))
		}
	}

	return make_type(resolved^)
}

// -- Type resolution helpers --

@(private = "file")
resolve_type_expr :: proc(s: ^Sema, te: Type_Expr) -> ^Resolved_Type {
	switch t in te {
	case Type_Named:
		sym := scope_lookup(s.global_scope, t.name)
		if sym == nil {
			sema_error(s, t.span, fmt.aprintf("undefined type '%s'", t.name))
			suggest_similar(s, t.name, t.span)
			return nil
		}
		return sym.type
	case Type_Array:
		elem := resolve_type_expr(s, t.elem^)
		size := 0
		if t.size != nil {
			size = eval_const_int(t.size)
			if size <= 0 {
				sema_error(s, t.span, "array size must be a positive integer constant")
				size = 0
			}
		}
		return make_type(Type_Array_Resolved{elem = elem, size = size})
	case Type_Tuple:
		// Should have been desugared before sema — report error
		sema_error(s, t.span, "tuple type not desugared (internal error)")
		return nil
	}
	return nil
}

@(private = "file")
resolve_arithmetic_type :: proc(s: ^Sema, left, right: ^Resolved_Type, span: Source_Span) -> ^Resolved_Type {
	// Same type -> same type
	if type_equals(left, right) do return left

	// scalar * vector -> vector  (and vice versa)
	if is_scalar(left) && is_vector(right) do return right
	if is_vector(left) && is_scalar(right) do return left

	// matrix * vector -> vector
	if is_matrix(left) && is_vector(right) do return right
	// vector * matrix -> vector
	if is_vector(left) && is_matrix(right) do return left

	// matrix * matrix -> matrix
	if is_matrix(left) && is_matrix(right) do return left

	// Mixed numeric scalar coercion
	if is_numeric_scalar(left) && is_numeric_scalar(right) {
		// If either is float/half, promote to float
		if is_float_scalar(left) do return left
		if is_float_scalar(right) do return right
		// Both integers: if either is uint, result is uint; otherwise int
		if is_uint_scalar(left) || is_uint_scalar(right) do return TYPE_UINT
		return TYPE_INT
	}

	return left
}

@(private = "file")
can_coerce :: proc(from, to: ^Resolved_Type) -> bool {
	if from == nil || to == nil do return false
	// Allow int -> float coercion
	if is_numeric_scalar(from) && is_numeric_scalar(to) do return true
	return false
}

is_valid_swizzle :: proc(field: string, vec_size: int) -> bool {
	if len(field) == 0 || len(field) > 4 do return false

	// xyzw set
	xyzw := "xyzw"
	// rgba set
	rgba := "rgba"
	// stpq set
	stpq := "stpq"

	in_xyzw := true
	in_rgba := true
	in_stpq := true

	for ch in field {
		idx_x := strings.index_byte(xyzw, u8(ch))
		idx_r := strings.index_byte(rgba, u8(ch))
		idx_s := strings.index_byte(stpq, u8(ch))

		if idx_x < 0 || idx_x >= vec_size do in_xyzw = false
		if idx_r < 0 || idx_r >= vec_size do in_rgba = false
		if idx_s < 0 || idx_s >= vec_size do in_stpq = false
	}

	return in_xyzw || in_rgba || in_stpq
}

@(private = "file")
sema_error :: proc(s: ^Sema, span: Source_Span, msg: string) {
	if s.error_count >= MAX_ERRORS {
		if s.error_count == MAX_ERRORS {
			append(&s.diagnostics, Diagnostic{
				level   = .Error,
				message = "too many errors, stopping",
				span    = span,
			})
			s.error_count += 1
		}
		return
	}
	append(&s.diagnostics, Diagnostic{
		level   = .Error,
		message = msg,
		span    = span,
	})
	s.error_count += 1
}

@(private = "file")
sema_note :: proc(s: ^Sema, span: Source_Span, msg: string) {
	if s.error_count > MAX_ERRORS do return
	append(&s.diagnostics, Diagnostic{
		level   = .Note,
		message = msg,
		span    = span,
	})
}

// "Did you mean?" suggestion — scan all symbols in scope chain
@(private = "file")
suggest_similar :: proc(s: ^Sema, name: string, span: Source_Span) {
	if len(name) < 3 do return

	best_name := ""
	best_dist := 3 // max distance to suggest

	scope := s.current_scope
	for scope != nil {
		for sym_name, _ in scope.symbols {
			d := levenshtein_distance(name, sym_name)
			if d < best_dist {
				best_dist = d
				best_name = sym_name
			}
		}
		scope = scope.parent
	}

	if best_name != "" {
		sema_note(s, span, fmt.aprintf("did you mean '%s'?", best_name))
	}
}

// Evaluate a constant integer expression (for array sizes).
// Supports integer literals and basic arithmetic (+, -, *).
@(private = "file")
eval_const_int :: proc(node: ^Ast_Node) -> int {
	if node == nil do return 0
	#partial switch node.kind {
	case .Literal:
		lit := node.derived.(^Ast_Literal)
		#partial switch v in lit.value {
		case i64: return int(v)
		case f64: return int(v)
		}
	case .Binary:
		bin := node.derived.(^Ast_Binary)
		l := eval_const_int(bin.left)
		r := eval_const_int(bin.right)
		#partial switch bin.op {
		case .Add: return l + r
		case .Sub: return l - r
		case .Mul: return l * r
		case .Div:
			if r != 0 do return l / r
		}
	case .Unary:
		un := node.derived.(^Ast_Unary)
		v := eval_const_int(un.operand)
		if un.op == .Neg do return -v
		return v
	}
	return 0
}