package shader

import "core:fmt"

// AST node types for the Luma shader language

Ast_Module :: struct {
	structs:     [dynamic]^Ast_Struct,
	functions:   [dynamic]^Ast_Function,
	bindings:    [dynamic]^Ast_Binding,
	constants:   [dynamic]^Ast_Const,
	shared_vars: [dynamic]^Ast_Shared,
}

Ast_Function :: struct {
	name:        string,
	params:      []Ast_Param,
	return_type: ^Type_Expr,
	body:        [dynamic]^Ast_Node,
	attributes:  []Ast_Attribute,
	span:        Source_Span,
}

Ast_Param :: struct {
	name:       string,
	type:       Type_Expr,
	attributes: []Ast_Attribute,
	span:       Source_Span,
}

Ast_Struct :: struct {
	name:       string,
	fields:     []Ast_Struct_Field,
	attributes: []Ast_Attribute,
	span:       Source_Span,
}

Ast_Struct_Field :: struct {
	name:       string,
	type:       Type_Expr,
	attributes: []Ast_Attribute,
	span:       Source_Span,
}

Ast_Binding :: struct {
	kind:       Binding_Kind,
	name:       string,
	type_expr:  Type_Expr,
	attributes: []Ast_Attribute,
	qualifier:  string,  // "storage" for buffer storage
	span:       Source_Span,
}

Binding_Kind :: enum {
	Uniform,
	Buffer,
}

Ast_Const :: struct {
	name:       string,
	type:       ^Type_Expr,
	value:      ^Ast_Node,
	attributes: []Ast_Attribute,
	span:       Source_Span,
}

Ast_Shared :: struct {
	name:      string,
	type_expr: Type_Expr,
	span:      Source_Span,
}

Ast_Discard :: struct {
	span: Source_Span,
}

Ast_Break :: struct {
	span: Source_Span,
}

Ast_Continue :: struct {
	span: Source_Span,
}

Ast_Block :: struct {
	stmts: [dynamic]^Ast_Node,
	span:  Source_Span,
}

Ast_Let :: struct {
	name:      string,
	type_expr: ^Type_Expr, // nil if inferred
	value:     ^Ast_Node,
	span:      Source_Span,
	mutable:   bool, // true for 'var', false for 'let'
}

Ast_Assign :: struct {
	target: ^Ast_Node,
	value:  ^Ast_Node,
	span:   Source_Span,
}

Ast_Output_Assign :: struct {
	name:     string,
	io_index: int,
	type:     Type_Expr,
	value:    ^Ast_Node,
	span:     Source_Span,
}

Ast_Return :: struct {
	value: ^Ast_Node, // nil for bare return
	span:  Source_Span,
}

Ast_If :: struct {
	condition:    ^Ast_Node,
	then_body:    [dynamic]^Ast_Node,
	elseif_clauses: []Ast_Elseif,
	else_body:    [dynamic]^Ast_Node, // empty if no else
	span:         Source_Span,
}

Ast_Elseif :: struct {
	condition: ^Ast_Node,
	body:      [dynamic]^Ast_Node,
}

Ast_For :: struct {
	var_name: string,
	start:    ^Ast_Node,
	stop:     ^Ast_Node,
	step:     ^Ast_Node, // nil if default step of 1
	body:     [dynamic]^Ast_Node,
	span:     Source_Span,
}

Ast_While :: struct {
	condition: ^Ast_Node,
	body:      [dynamic]^Ast_Node,
	span:      Source_Span,
}

Ast_Call :: struct {
	callee: ^Ast_Node,
	args:   []^Ast_Node,
	span:   Source_Span,
}

Ast_Field_Access :: struct {
	object: ^Ast_Node,
	field:  string,
	span:   Source_Span,
}

Ast_Index :: struct {
	object: ^Ast_Node,
	index:  ^Ast_Node,
	span:   Source_Span,
}

Ast_Swizzle :: struct {
	object:     ^Ast_Node,
	components: string, // e.g. "xyz", "rg"
	span:       Source_Span,
}

Ast_Unary :: struct {
	op:      Unary_Op,
	operand: ^Ast_Node,
	span:    Source_Span,
}

Unary_Op :: enum {
	Neg,
	Not,
}

Ast_Binary :: struct {
	op:    Binary_Op,
	left:  ^Ast_Node,
	right: ^Ast_Node,
	span:  Source_Span,
}

Binary_Op :: enum {
	Add,
	Sub,
	Mul,
	Div,
	Mod,
	Eq,
	Neq,
	Lt,
	Gt,
	Lte,
	Gte,
	And,
	Or,
}

Ast_Literal :: struct {
	value: Literal_Value,
	span:  Source_Span,
}

Literal_Value :: union {
	i64,
	f64,
	bool,
	string,
}

Ast_Ident :: struct {
	name: string,
	span: Source_Span,
}

Ast_Struct_Literal :: struct {
	type_name: string,
	fields:    []Ast_Struct_Literal_Field,
	span:      Source_Span,
}

Ast_Struct_Literal_Field :: struct {
	name:  string,
	value: ^Ast_Node,
	span:  Source_Span,
}

Ast_Attribute :: struct {
	name: string,
	args: []string,
	span: Source_Span,
}

// Type expression (used in the AST before resolution)

Type_Expr :: union {
	Type_Named,
	Type_Array,
	Type_Tuple, // (name: Type, ...) — inline entry point return types
}

Type_Named :: struct {
	name: string,
	span: Source_Span,
}

Type_Array :: struct {
	elem: ^Type_Expr,
	size: ^Ast_Node, // nil for runtime-sized ([]T)
	span: Source_Span,
}

Type_Tuple :: struct {
	fields: []Ast_Struct_Field, // (name: Type, ...)
	span:   Source_Span,
}

// The main AST node union

Ast_Node :: struct {
	kind:          Ast_Node_Kind,
	resolved_type: ^Resolved_Type, // filled in by sema
	span:          Source_Span,
	derived:       Ast_Node_Derived,
}

Ast_Node_Derived :: union {
	^Ast_Function,
	^Ast_Struct,
	^Ast_Binding,
	^Ast_Block,
	^Ast_Let,
	^Ast_Assign,
	^Ast_Output_Assign,
	^Ast_Return,
	^Ast_If,
	^Ast_For,
	^Ast_While,
	^Ast_Call,
	^Ast_Field_Access,
	^Ast_Index,
	^Ast_Swizzle,
	^Ast_Unary,
	^Ast_Binary,
	^Ast_Literal,
	^Ast_Ident,
	^Ast_Struct_Literal,
	^Ast_Discard,
	^Ast_Break,
	^Ast_Continue,
}

Ast_Node_Kind :: enum {
	Function,
	Struct,
	Binding,
	Block,
	Let,
	Assign,
	Output_Assign,
	Return,
	If,
	For,
	While,
	Call,
	Field_Access,
	Index,
	Swizzle,
	Unary,
	Binary,
	Literal,
	Ident,
	Struct_Literal,
	Discard,
	Break,
	Continue,
}

// Helper to create AST nodes

make_node :: proc(kind: Ast_Node_Kind, derived: Ast_Node_Derived, span: Source_Span, allocator := context.allocator) -> ^Ast_Node {
	node := new(Ast_Node, allocator)
	node.kind = kind
	node.derived = derived
	node.span = span
	return node
}

// -- AST Printer --

ast_to_string :: proc(module: ^Ast_Module, allocator := context.allocator) -> string {
	w := writer_init(allocator)

	write_line(&w, "=== AST Module ===")
	write_line(&w, "")

	for c in module.constants {
		ast_print_const(&w, c)
	}
	if len(module.constants) > 0 do write_line(&w, "")

	for s in module.structs {
		ast_print_struct(&w, s)
		write_line(&w, "")
	}

	for b in module.bindings {
		ast_print_binding(&w, b)
	}
	if len(module.bindings) > 0 do write_line(&w, "")

	for sv in module.shared_vars {
		write_line(&w, "shared ", sv.name, ": ", type_expr_str(sv.type_expr))
	}
	if len(module.shared_vars) > 0 do write_line(&w, "")

	for f in module.functions {
		ast_print_function(&w, f)
		write_line(&w, "")
	}

	return writer_to_string(w)
}

@(private = "file")
ast_print_attrs :: proc(w: ^Writer, attrs: []Ast_Attribute) {
	for attr in attrs {
		write(w, "@", attr.name)
		if len(attr.args) > 0 {
			write(w, "(")
			for a, i in attr.args {
				if i > 0 do write(w, ", ")
				write(w, a)
			}
			write(w, ")")
		}
		write(w, " ")
	}
}

@(private = "file")
ast_print_const :: proc(w: ^Writer, c: ^Ast_Const) {
	for _ in 0 ..< w.indent do write(w, "\t")
	write(w, "const ", c.name)
	if c.type != nil {
		write(w, ": ", type_expr_str(c.type^))
	}
	if c.value != nil {
		write(w, " = ")
		ast_print_node(w, c.value)
	}
	write(w, "\n")
}

@(private = "file")
ast_print_struct :: proc(w: ^Writer, s: ^Ast_Struct) {
	for _ in 0 ..< w.indent do write(w, "\t")
	ast_print_attrs(w, s.attributes)
	write(w, "struct ", s.name, "\n")
	indent(w)
	for f in s.fields {
		for _ in 0 ..< w.indent do write(w, "\t")
		ast_print_attrs(w, f.attributes)
		write(w, f.name, ": ", type_expr_str(f.type), "\n")
	}
	dedent(w)
	write_line(w, "end")
}

@(private = "file")
ast_print_binding :: proc(w: ^Writer, b: ^Ast_Binding) {
	for _ in 0 ..< w.indent do write(w, "\t")
	ast_print_attrs(w, b.attributes)
	kind_str := b.kind == .Uniform ? "uniform" : "buffer"
	if b.qualifier != "" {
		write(w, kind_str, " ", b.qualifier, " ", b.name, ": ", type_expr_str(b.type_expr), "\n")
	} else {
		write(w, kind_str, " ", b.name, ": ", type_expr_str(b.type_expr), "\n")
	}
}

@(private = "file")
ast_print_function :: proc(w: ^Writer, f: ^Ast_Function) {
	for _ in 0 ..< w.indent do write(w, "\t")
	ast_print_attrs(w, f.attributes)
	write(w, "function ", f.name, "(")
	for p, i in f.params {
		if i > 0 do write(w, ", ")
		ast_print_attrs(w, p.attributes)
		write(w, p.name, ": ", type_expr_str(p.type))
	}
	write(w, ")")
	if f.return_type != nil {
		write(w, " -> ", type_expr_str(f.return_type^))
	}
	write(w, "\n")
	indent(w)
	for node in f.body {
		ast_print_node(w, node)
	}
	dedent(w)
	write_line(w, "end")
}

@(private = "file")
ast_print_node :: proc(w: ^Writer, node: ^Ast_Node) {
	if node == nil {
		write(w, "<nil>")
		return
	}

	switch d in node.derived {
	case ^Ast_Function:
		ast_print_function(w, d)
	case ^Ast_Struct:
		ast_print_struct(w, d)
	case ^Ast_Binding:
		ast_print_binding(w, d)
	case ^Ast_Block:
		for s in d.stmts {
			ast_print_node(w, s)
		}
	case ^Ast_Let:
		for _ in 0 ..< w.indent do write(w, "\t")
		write(w, "let ", d.name)
		if d.type_expr != nil {
			write(w, ": ", type_expr_str(d.type_expr^))
		}
		if d.value != nil {
			write(w, " = ")
			ast_print_node(w, d.value)
		}
		write(w, "\n")
	case ^Ast_Assign:
		for _ in 0 ..< w.indent do write(w, "\t")
		ast_print_node(w, d.target)
		write(w, " = ")
		ast_print_node(w, d.value)
		write(w, "\n")
	case ^Ast_Output_Assign:
		for _ in 0 ..< w.indent do write(w, "\t")
		write(w, "out.", d.name, " = ")
		ast_print_node(w, d.value)
		write(w, "\n")
	case ^Ast_Return:
		for _ in 0 ..< w.indent do write(w, "\t")
		write(w, "return")
		if d.value != nil {
			write(w, " ")
			ast_print_node(w, d.value)
		}
		write(w, "\n")
	case ^Ast_If:
		for _ in 0 ..< w.indent do write(w, "\t")
		write(w, "if ")
		ast_print_node(w, d.condition)
		write(w, " then\n")
		indent(w)
		for s in d.then_body {
			ast_print_node(w, s)
		}
		dedent(w)
		for ei in d.elseif_clauses {
			for _ in 0 ..< w.indent do write(w, "\t")
			write(w, "elseif ")
			ast_print_node(w, ei.condition)
			write(w, " then\n")
			indent(w)
			for s in ei.body {
				ast_print_node(w, s)
			}
			dedent(w)
		}
		if len(d.else_body) > 0 {
			write_line(w, "else")
			indent(w)
			for s in d.else_body {
				ast_print_node(w, s)
			}
			dedent(w)
		}
		write_line(w, "end")
	case ^Ast_For:
		for _ in 0 ..< w.indent do write(w, "\t")
		write(w, "for ", d.var_name, " = ")
		ast_print_node(w, d.start)
		write(w, ", ")
		ast_print_node(w, d.stop)
		if d.step != nil {
			write(w, ", ")
			ast_print_node(w, d.step)
		}
		write(w, " do\n")
		indent(w)
		for s in d.body {
			ast_print_node(w, s)
		}
		dedent(w)
		write_line(w, "end")
	case ^Ast_While:
		for _ in 0 ..< w.indent do write(w, "\t")
		write(w, "while ")
		ast_print_node(w, d.condition)
		write(w, " do\n")
		indent(w)
		for s in d.body {
			ast_print_node(w, s)
		}
		dedent(w)
		write_line(w, "end")
	case ^Ast_Call:
		ast_print_node(w, d.callee)
		write(w, "(")
		for arg, i in d.args {
			if i > 0 do write(w, ", ")
			ast_print_node(w, arg)
		}
		write(w, ")")
	case ^Ast_Field_Access:
		ast_print_node(w, d.object)
		write(w, ".", d.field)
	case ^Ast_Index:
		ast_print_node(w, d.object)
		write(w, "[")
		ast_print_node(w, d.index)
		write(w, "]")
	case ^Ast_Swizzle:
		ast_print_node(w, d.object)
		write(w, ".", d.components)
	case ^Ast_Unary:
		op_str := d.op == .Neg ? "-" : "not "
		write(w, "(", op_str)
		ast_print_node(w, d.operand)
		write(w, ")")
	case ^Ast_Binary:
		write(w, "(")
		ast_print_node(w, d.left)
		write(w, " ", binary_op_str(d.op), " ")
		ast_print_node(w, d.right)
		write(w, ")")
	case ^Ast_Literal:
		switch v in d.value {
		case i64:    write(w, v)
		case f64:    write(w, v)
		case bool:   write(w, v ? "true" : "false")
		case string: write(w, "\"", v, "\"")
		}
	case ^Ast_Ident:
		write(w, d.name)
	case ^Ast_Struct_Literal:
		write(w, d.type_name, " { ")
		for f, i in d.fields {
			if i > 0 do write(w, ", ")
			write(w, f.name, " = ")
			ast_print_node(w, f.value)
		}
		write(w, " }")
	case ^Ast_Discard:
		for _ in 0 ..< w.indent do write(w, "\t")
		write(w, "discard\n")
	case ^Ast_Break:
		for _ in 0 ..< w.indent do write(w, "\t")
		write(w, "break\n")
	case ^Ast_Continue:
		for _ in 0 ..< w.indent do write(w, "\t")
		write(w, "continue\n")
	}
}

@(private = "file")
type_expr_str :: proc(te: Type_Expr) -> string {
	switch t in te {
	case Type_Named:
		return t.name
	case Type_Array:
		if t.elem != nil {
			return fmt.aprintf("[]%s", type_expr_str(t.elem^))
		}
		return "[]?"
	case Type_Tuple:
		w := writer_init()
		write(&w, "(")
		for f, i in t.fields {
			if i > 0 do write(&w, ", ")
			write(&w, f.name, ": ", type_expr_str(f.type))
		}
		write(&w, ")")
		return writer_to_string(w)
	}
	return "?"
}

@(private = "file")
binary_op_str :: proc(op: Binary_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 "and"
	case .Or:  return "or"
	}
	return "?"
}