package shader

import "core:os"
import "core:strings"
import "core:path/filepath"

MAX_INCLUDE_DEPTH :: 32

// Source map entry: maps a range in the flattened source back to original file + line
Source_Map_Entry :: struct {
	flat_line_start: int,    // line in flattened source
	flat_line_count: int,    // number of lines from this file segment
	orig_file:       string, // original file path
	orig_line_start: int,    // starting line in original file
}

Source_Map :: struct {
	entries: [dynamic]Source_Map_Entry,
}

Preprocessor :: struct {
	include_dirs: []string,
	file_reader:  File_Reader,
	included:     map[string]bool,  // for #pragma once (by resolved path)
	diagnostics:  [dynamic]Diagnostic,
	source_map:   Source_Map,
}

File_Reader :: #type proc(path: string) -> (content: string, ok: bool)

default_file_reader :: proc(path: string) -> (content: string, ok: bool) {
	data, err := os.read_entire_file(path, context.allocator)
	if err != nil {
		return "", false
	}
	return string(data), true
}

preprocessor_init :: proc(include_dirs: []string = nil, reader: File_Reader = nil, allocator := context.allocator) -> Preprocessor {
	return Preprocessor{
		include_dirs = include_dirs,
		file_reader  = reader != nil ? reader : default_file_reader,
		included     = make(map[string]bool, allocator = allocator),
		diagnostics  = make([dynamic]Diagnostic, allocator),
		source_map   = Source_Map{entries = make([dynamic]Source_Map_Entry, allocator)},
	}
}

// Preprocess expands #include directives and returns flattened source
preprocess :: proc(pp: ^Preprocessor, source: string, file: string) -> (string, []Diagnostic) {
	result := preprocess_recursive(pp, source, file, 0)
	return result, pp.diagnostics[:]
}

@(private = "file")
preprocess_recursive :: proc(pp: ^Preprocessor, source: string, file: string, depth: int) -> string {
	if depth > MAX_INCLUDE_DEPTH {
		append(&pp.diagnostics, Diagnostic{
			level   = .Error,
			message = "maximum include depth exceeded",
			span    = Source_Span{file = file, line_start = 1, col_start = 1},
		})
		return source
	}

	buf := strings.builder_make()
	lines := strings.split(source, "\n")
	current_flat_line := 1
	segment_start_line := 1

	for i := 0; i < len(lines); i += 1 {
		line := strings.trim_space(lines[i])

		if strings.has_prefix(line, "#pragma once") {
			abs_path := resolve_path(file)
			if abs_path in pp.included {
				return "" // already included
			}
			pp.included[abs_path] = true
			continue
		}

		if strings.has_prefix(line, "#include") {
			// Flush source map entry for lines before this include
			if i > segment_start_line - 1 {
				append(&pp.source_map.entries, Source_Map_Entry{
					flat_line_start = current_flat_line,
					flat_line_count = i - (segment_start_line - 1),
					orig_file       = file,
					orig_line_start = segment_start_line,
				})
				current_flat_line += i - (segment_start_line - 1)
			}

			include_path, ok := parse_include_directive(line)
			if !ok {
				append(&pp.diagnostics, Diagnostic{
					level   = .Error,
					message = "malformed #include directive",
					span    = Source_Span{file = file, line_start = i + 1, col_start = 1},
				})
				continue
			}

			resolved, is_angle := resolve_include_path(pp, include_path, file)

			// If filesystem resolution failed, try the file_reader directly
			// (supports embedded/virtual includes)
			if resolved == "" {
				if _, reader_ok := pp.file_reader(include_path); reader_ok {
					resolved = include_path
				}
			}

			if resolved == "" {
				append(&pp.diagnostics, Diagnostic{
					level   = .Error,
					message = strings.concatenate({"could not find included file: ", include_path}),
					span    = Source_Span{file = file, line_start = i + 1, col_start = 1},
				})
				continue
			}

			// Check circular include
			abs_resolved := resolve_path(resolved)
			if abs_resolved in pp.included {
				// Already included via #pragma once — skip silently
				segment_start_line = i + 2
				continue
			}

			content, read_ok := pp.file_reader(resolved)
			if !read_ok {
				append(&pp.diagnostics, Diagnostic{
					level   = .Error,
					message = strings.concatenate({"could not read included file: ", resolved}),
					span    = Source_Span{file = file, line_start = i + 1, col_start = 1},
				})
				continue
			}

			expanded := preprocess_recursive(pp, content, resolved, depth + 1)
			strings.write_string(&buf, expanded)
			if len(expanded) > 0 && expanded[len(expanded) - 1] != '\n' {
				strings.write_string(&buf, "\n")
			}

			segment_start_line = i + 2
			continue
		}

		strings.write_string(&buf, lines[i])
		if i < len(lines) - 1 {
			strings.write_string(&buf, "\n")
		}
	}

	return strings.to_string(buf)
}

@(private = "file")
parse_include_directive :: proc(line: string) -> (path: string, ok: bool) {
	rest := strings.trim_space(line[len("#include"):])
	if len(rest) < 2 do return "", false

	if rest[0] == '"' {
		end := strings.index_byte(rest[1:], '"')
		if end < 0 do return "", false
		return rest[1:1 + end], true
	}
	if rest[0] == '<' {
		end := strings.index_byte(rest[1:], '>')
		if end < 0 do return "", false
		return rest[1:1 + end], true
	}
	return "", false
}

@(private = "file")
resolve_include_path :: proc(pp: ^Preprocessor, path: string, current_file: string) -> (resolved: string, is_angle: bool) {
	// Try relative to current file first
	dir := filepath.dir(current_file)
	candidate, join_err := filepath.join({dir, path}, context.allocator)
	if join_err == nil && os.exists(candidate) {
		return candidate, false
	}

	// Try include dirs
	for inc_dir in pp.include_dirs {
		c, jerr := filepath.join({inc_dir, path}, context.allocator)
		if jerr == nil && os.exists(c) {
			return c, true
		}
	}

	return "", false
}

@(private = "file")
resolve_path :: proc(path: string) -> string {
	abs, err := filepath.abs(path, context.allocator)
	if err == nil {
		return abs
	}
	return path
}

// Look up original file and line from a flattened source position
source_map_lookup :: proc(sm: ^Source_Map, flat_line: int) -> (file: string, orig_line: int) {
	for entry in sm.entries {
		if flat_line >= entry.flat_line_start && flat_line < entry.flat_line_start + entry.flat_line_count {
			offset := flat_line - entry.flat_line_start
			return entry.orig_file, entry.orig_line_start + offset
		}
	}
	return "", flat_line
}