package loader

import "../ast"
import "../cimport"
import "../lexer"
import "../parser"
import "../source"
import "../symbol"
import "../target"
import "../types"
import "core:fmt"
import "core:mem"
import "core:os"
import "core:path/filepath"
import "core:slice"
import "core:strings"

State :: struct {
	module:          ^ast.Module,
	sources:         ^source.Store,
	diagnostics:     ^source.Diagnostics,
	symbols:         ^symbol.Table,
	token_allocator: mem.Allocator,
	allocator:       mem.Allocator,
	c_options:       cimport.Options,
	selected:        target.Target,
	record_identities: [dynamic]string,
	record_types:      [dynamic]types.Type,
	root_failed:     bool,
}

is_identifier :: proc(value: string) -> bool {
	if len(value) == 0 {
		return false
	}
	is_start := proc(value: byte) -> bool {
		return value == '_' || value >= 'a' && value <= 'z' || value >= 'A' && value <= 'Z'
	}
	if !is_start(value[0]) {
		return false
	}
	for byte_value in transmute([]byte)value[1:] {
		if !is_start(byte_value) && !(byte_value >= '0' && byte_value <= '9') {
			return false
		}
	}
	return true
}

find_package :: proc(state: ^State, path: string) -> ast.Package_Id {
	for pkg, id in state.module.packages {
		if pkg.path == path {
			return ast.package_id(id)
		}
	}
	return ast.INVALID_PACKAGE
}

add_placeholder :: proc(state: ^State, path: string) -> ast.Package_Id {
	if existing := find_package(state, path); existing != ast.INVALID_PACKAGE {
		return existing
	}
	id := ast.package_id(len(state.module.packages))
	append(&state.module.packages, ast.Package{
		path=strings.clone(path, state.allocator),
		name=symbol.intern(state.symbols, filepath.base(path)),
		available=false,
	})
	return id
}

read_package_files :: proc(state: ^State, path: string) -> ([]os.File_Info, bool) {
	handle, open_error := os.open(path, os.O_RDONLY)
	if open_error != nil {
		return nil, false
	}
	defer os.close(handle)
	entries, read_error := os.read_dir(handle, -1, state.allocator)
	if read_error != nil {
		return nil, false
	}
	slice.sort_by(entries, proc(a, b: os.File_Info) -> bool {
		return a.name < b.name
	})
	files: [dynamic]os.File_Info
	files.allocator = state.allocator
	for entry in entries {
		if !entry.is_dir && filepath.ext(entry.name) == ".bro" {
			append(&files, entry)
		} else {
			os.file_info_delete(entry, state.allocator)
		}
	}
	delete(entries, state.allocator)
	return files[:], true
}

resolve_import_path :: proc(state: ^State, importing_path, import_path: string) -> (string, bool) {
	if filepath.is_abs(import_path) {
		return "", false
	}
	joined, join_error := filepath.join({importing_path, import_path}, state.allocator)
	if join_error != nil {
		return "", false
	}
	canonical, ok := filepath.abs(joined, state.allocator)
	if ok {
		delete(joined, state.allocator)
		return canonical, true
	}
	return joined, false
}

header_package_name :: proc(path: string, symbols: ^symbol.Table) -> symbol.Id {
	base := filepath.base(path)
	extension := filepath.ext(base)
	if len(extension) > 0 {
		base = base[:len(base)-len(extension)]
	}
	return symbol.intern(symbols, base)
}

find_record_identity :: proc(state: ^State, identity: string) -> types.Type {
	for existing, index in state.record_identities {
		if existing == identity {
			return state.record_types[index]
		}
	}
	return types.INVALID
}

translate_c_type :: proc(
	state: ^State,
	result: ^cimport.Result,
	value: cimport.Type_Id,
	pkg: ast.Package_Id,
	record_mapping: []types.Type,
	type_mapping: []types.Type,
) -> types.Type {
	if value == cimport.INVALID_TYPE || int(value) < 0 || int(value) >= len(result.types) {
		return types.INVALID
	}
	if types.is_valid(type_mapping[value]) {
		return type_mapping[value]
	}
	item := result.types[value]
	translated := types.INVALID
	switch item.kind {
	case .Invalid:      translated = types.INVALID
	case .Void:         translated = types.VOID
	case .C_Char:       translated = types.C_CHAR
	case .C_Schar:      translated = types.C_SCHAR
	case .C_Uchar:      translated = types.C_UCHAR
	case .C_Short:      translated = types.C_SHORT
	case .C_Ushort:     translated = types.C_USHORT
	case .C_Int:        translated = types.C_INT
	case .C_Uint:       translated = types.C_UINT
	case .C_Long:       translated = types.C_LONG
	case .C_Ulong:      translated = types.C_ULONG
	case .C_Longlong:   translated = types.C_LONGLONG
	case .C_Ulonglong:  translated = types.C_ULONGLONG
	case .C_Float:      translated = types.C_FLOAT
	case .C_Double:     translated = types.C_DOUBLE
	case .C_Longdouble: translated = types.C_LONGDOUBLE
	case .Pointer:
		child := translate_c_type(state, result, item.child, pkg, record_mapping, type_mapping)
		if types.is_valid(child) {
			pointer := types.pointer(&state.module.type_store, child, item.mutable, true)
			translated = types.optional(&state.module.type_store, pointer)
		}
	case .Array:
		child := translate_c_type(state, result, item.child, pkg, record_mapping, type_mapping)
		if types.is_valid(child) {
			translated = types.array(&state.module.type_store, child, item.count, item.mutable)
		}
	case .Function:
		params := make([]types.Type, len(item.params), state.allocator)
		for param, index in item.params {
			params[index] = translate_c_type(state, result, param, pkg, record_mapping, type_mapping)
			if !types.is_valid(params[index]) {
				delete(params, state.allocator)
				type_mapping[value] = types.INVALID
				return types.INVALID
			}
		}
		result_type := translate_c_type(state, result, item.child, pkg, record_mapping, type_mapping)
		if types.is_valid(result_type) {
			translated = types.function(&state.module.type_store, params, result_type, true, item.variadic)
		}
		delete(params, state.allocator)
	case .Record:
		if int(item.record) < len(record_mapping) {
			translated = record_mapping[item.record]
		}
	}
	type_mapping[value] = translated
	return translated
}

function_signatures_equal :: proc(left: ast.Function, params: []ast.Param, result: types.Type, variadic: bool) -> bool {
	if left.result != result || left.variadic != variadic || len(left.params) != len(params) {
		return false
	}
	for param, index in params {
		if left.params[index].type != param.type {
			return false
		}
	}
	return true
}

RECORD_DEPENDENCY_PENDING :: "C record contains an incomplete or unsupported record field"

record_layout_reason :: proc(
	store: ^types.Store,
	record: cimport.Record,
	fields: []types.Field,
	selected: target.Target,
) -> string {
	if len(fields) == 0 {
		if record.size > 0 {
			return "anonymous C record fields are not supported"
		}
		return "empty C records are not supported"
	}
	if len(fields) != len(record.fields) || record.alignment == 0 {
		return "C record metadata is incomplete"
	}
	size: u64
	alignment: u64 = 1
	if record.kind == .Union {
		for field in fields {
			if !types.is_runtime_value(field.type, store) {
				return RECORD_DEPENDENCY_PENDING
			}
			if field.offset != 0 {
				return "non-natural C record layouts are not supported"
			}
			size = max(size, types.size(field.type, store, selected))
			alignment = max(alignment, u64(types.alignment_of(field.type, store, selected)))
		}
	} else {
		offset: u64
		for field in fields {
			if !types.is_runtime_value(field.type, store) {
				return RECORD_DEPENDENCY_PENDING
			}
			field_alignment := u64(types.alignment_of(field.type, store, selected))
			offset = (offset+field_alignment-1)/field_alignment*field_alignment
			if field.offset != offset {
				if field.offset < offset {
					return "packed C records are not supported"
				}
				return "non-natural C record layouts are not supported"
			}
			offset += types.size(field.type, store, selected)
			alignment = max(alignment, field_alignment)
		}
		size = offset
	}
	size = (size+alignment-1)/alignment*alignment
	if u64(record.alignment) > alignment {
		return "over-aligned C records are not supported"
	}
	if u64(record.alignment) < alignment || record.size < size {
		return "packed C records are not supported"
	}
	if record.size != size {
		return "non-natural C record layouts are not supported"
	}
	return ""
}

c_record_by_value_reason :: proc(result: ^cimport.Result, value: cimport.Type_Id, depth := 0) -> string {
	if depth > 64 || value == cimport.INVALID_TYPE || int(value) < 0 || int(value) >= len(result.types) {
		return ""
	}
	item := result.types[value]
	#partial switch item.kind {
	case .Pointer:
		return ""
	case .Array:
		return c_record_by_value_reason(result, item.child, depth+1)
	case .Record:
		if int(item.record) >= len(result.records) {
			return "C record metadata is incomplete"
		}
		record := result.records[item.record]
		if len(record.reason) > 0 {
			return record.reason
		}
		if !record.complete {
			return "incomplete C records are pointer-only"
		}
	}
	return ""
}

load_header :: proc(state: ^State, path: string, import_span: source.Span) -> ast.Package_Id {
	canonical, ok := filepath.abs(path, state.allocator)
	if !ok {
		id := add_placeholder(state, path)
		source.addf(state.diagnostics, import_span, "could not resolve C header '%s'", path)
		return id
	}
	if existing := find_package(state, canonical); existing != ast.INVALID_PACKAGE {
		delete(canonical, state.allocator)
		return existing
	}
	pkg_id := ast.package_id(len(state.module.packages))
	append(&state.module.packages, ast.Package{
		path=canonical,
		name=header_package_name(canonical, state.symbols),
		available=false,
		kind=.C_Header,
	})
	result := cimport.import_header(state.c_options, canonical, state.selected, state.allocator)
	defer cimport.destroy_result(&result)
	if !result.available {
		message := result.error_message if len(result.error_message) > 0 else "C header import failed"
		source.addf(state.diagnostics, import_span, "could not import C header '%s': %s", path, message)
		if result.infrastructure {
			state.root_failed = true
		}
		return pkg_id
	}
	state.module.packages[pkg_id].available = true

	record_mapping := make([]types.Type, len(result.records), state.allocator)
	defer delete(record_mapping, state.allocator)
	for record, index in result.records {
		record_type := find_record_identity(state, record.identity)
		if !types.is_valid(record_type) {
			name := record.name
			if len(name) == 0 {
				name = fmt.tprintf("__c_record_%d", len(state.record_types))
			}
			record_type = types.named(&state.module.type_store, u32(pkg_id), u32(symbol.intern(state.symbols, name)))
			_ = types.define_record(&state.module.type_store, record_type, nil, true, true, record.kind == .Union)
			append(&state.record_identities, strings.clone(record.identity, state.allocator))
			append(&state.record_types, record_type)
		}
		record_mapping[index] = record_type
	}
	type_mapping := make([]types.Type, len(result.types), state.allocator)
	defer delete(type_mapping, state.allocator)

	for _ in 0..<len(result.records) {
		changed := false
		for record, record_index in result.records {
			if !record.complete || len(record.reason) > 0 {
				continue
			}
			record_type := record_mapping[record_index]
			item, ok := types.node(&state.module.type_store, record_type)
			if !ok || (item.declared && !item.opaque) {
				continue
			}
			fields := make([]types.Field, len(record.fields), state.allocator)
			for field, field_index in record.fields {
				fields[field_index] = types.Field{
					name=u32(symbol.intern(state.symbols, field.name)),
					type=translate_c_type(state, &result, field.type, pkg_id, record_mapping, type_mapping),
					offset=field.offset,
				}
			}
			layout_reason := record_layout_reason(&state.module.type_store, record, fields, state.selected)
			if len(layout_reason) == 0 {
				changed = types.define_record(
					&state.module.type_store, record_type, fields, true, false,
					record.kind == .Union, record.size, record.alignment,
				) || changed
			} else if layout_reason != RECORD_DEPENDENCY_PENDING && len(record.reason) == 0 {
				delete(result.records[record_index].reason, result.allocator)
				result.records[record_index].reason = fmt.aprintf("%s", layout_reason, allocator=result.allocator)
			}
			delete(fields, state.allocator)
		}
		if !changed {
			break
		}
	}

	for alias in result.aliases {
		name := symbol.intern(state.symbols, alias.name)
		id := types.named(&state.module.type_store, u32(pkg_id), u32(name))
		child := translate_c_type(state, &result, alias.type, pkg_id, record_mapping, type_mapping)
		_ = types.define_alias(&state.module.type_store, id, child)
		if len(alias.reason) > 0 {
			append(&state.module.unsupported, ast.Unsupported{
				pkg=pkg_id,
				name=name,
				reason=strings.clone(alias.reason, state.allocator),
			})
		}
	}

	for function in result.functions {
		params := make([]ast.Param, len(function.params), state.allocator)
		for param_type, index in function.params {
			params[index] = ast.Param{
				name=symbol.intern(state.symbols, fmt.tprintf("arg%d", index)),
				span=import_span,
				type=translate_c_type(state, &result, param_type, pkg_id, record_mapping, type_mapping),
			}
		}
		function_result := translate_c_type(state, &result, function.result, pkg_id, record_mapping, type_mapping)
		unsupported_reason := function.reason
		if len(unsupported_reason) == 0 {
			for param_type in function.params {
				if reason := c_record_by_value_reason(&result, param_type); len(reason) > 0 {
					unsupported_reason = reason
					break
				}
			}
		}
		if len(unsupported_reason) == 0 {
			unsupported_reason = c_record_by_value_reason(&result, function.result)
		}
		if len(unsupported_reason) == 0 {
			for param in params {
				if types.contains_c_struct_by_value(param.type, &state.module.type_store) {
					unsupported_reason = "C record parameter is incomplete or has an unsupported layout"
					break
				}
			}
		}
		if len(unsupported_reason) == 0 &&
		   types.contains_c_struct_by_value(function_result, &state.module.type_store) {
			unsupported_reason = "C record result is incomplete or has an unsupported layout"
		}
		name := symbol.intern(state.symbols, function.name)
		duplicate := false
		for &existing in state.module.functions {
			if existing.pkg != pkg_id || existing.name != name {
				continue
			}
			duplicate = true
			if !function_signatures_equal(existing, params, function_result, function.variadic) && len(existing.unsupported_reason) == 0 {
				existing.unsupported_reason = fmt.aprintf(
					"conflicting C declarations for '%s'",
					function.name,
					allocator=state.allocator,
				)
			}
			break
		}
		if duplicate {
			delete(params, state.allocator)
			continue
		}
		append(&state.module.functions, ast.Function{
			span=import_span,
			name=name,
			pkg=pkg_id,
			file=ast.INVALID_FILE,
			c_abi=true,
			imported=true,
			has_body=false,
			variadic=function.variadic,
			params=params,
			result=function_result,
			unsupported_reason=strings.clone(unsupported_reason, state.allocator),
			diagnostic=source.INVALID_DIAGNOSTIC,
		})
	}
	for item in result.unsupported {
		append(&state.module.unsupported, ast.Unsupported{
			pkg=pkg_id,
			name=symbol.intern(state.symbols, item.name),
			reason=strings.clone(item.reason, state.allocator),
		})
	}
	return pkg_id
}

load_package :: proc(state: ^State, path: string, import_span: source.Span, is_root := false) -> ast.Package_Id {
	canonical, ok := filepath.abs(path, state.allocator)
	if !ok || !os.is_dir(path) {
		if is_root {
			state.root_failed = true
			if len(canonical) > 0 {
				delete(canonical, state.allocator)
			}
			return ast.INVALID_PACKAGE
		}
		placeholder := path
		if len(canonical) > 0 {
			placeholder = canonical
		}
		id := add_placeholder(state, placeholder)
		source.addf(state.diagnostics, import_span, "could not import package directory '%s'", path)
		if len(canonical) > 0 {
			delete(canonical, state.allocator)
		}
		return id
	}
	if existing := find_package(state, canonical); existing != ast.INVALID_PACKAGE {
		delete(canonical, state.allocator)
		return existing
	}

	pkg_id := ast.package_id(len(state.module.packages))
	append(&state.module.packages, ast.Package{
		path=canonical,
		name=symbol.intern(state.symbols, filepath.base(canonical)),
		available=true,
	})
	files, files_ok := read_package_files(state, canonical)
	if !files_ok {
		state.root_failed = true
		return pkg_id
	}
	if len(files) == 0 {
		if is_root {
			state.root_failed = true
		} else {
			source.addf(state.diagnostics, import_span, "package '%s' contains no readable .bro files", canonical)
			state.module.packages[pkg_id].available = false
		}
		os.file_info_slice_delete(files, state.allocator)
		return pkg_id
	}

	for file_info in files {
		if file_info.size < 0 || !source.fits_source_length(u64(file_info.size)) {
			source.addf(state.diagnostics, import_span, "source file '%s' exceeds the 4 GiB source limit", file_info.fullpath)
			state.root_failed = true
			continue
		}
		bytes, read_ok := os.read_entire_file(file_info.fullpath, state.sources.allocator)
		if !read_ok {
			state.root_failed = true
			continue
		}
		if !source.fits_source_length(u64(len(bytes))) {
			source.addf(state.diagnostics, import_span, "source file '%s' exceeds the 4 GiB source limit", file_info.fullpath)
			delete(bytes, state.sources.allocator)
			state.root_failed = true
			continue
		}
		source_id := source.add_source_owned(state.sources, file_info.fullpath, bytes)
		file_id := ast.file_id(len(state.module.files))
		append(&state.module.files, ast.File{source=source_id, pkg=pkg_id})
		stream := lexer.lex(&state.sources.items[source_id], state.diagnostics, state.symbols, state.token_allocator)
		parser.parse_into(&stream, &state.sources.items[source_id], state.diagnostics, state.module, pkg_id, file_id)
		delete(stream.items)
	}
	os.file_info_slice_delete(files, state.allocator)

	import_count := len(state.module.imports)
	for import_id in 0..<import_count {
		import_item := state.module.imports[import_id]
		if import_item.pkg != pkg_id || import_item.target != ast.INVALID_PACKAGE {
			continue
		}
		if filepath.is_abs(import_item.path) {
			state.module.imports[import_id].diagnostic = source.add(state.diagnostics, import_item.span, "absolute import paths are invalid")
			state.module.imports[import_id].valid = false
			state.module.imports[import_id].target = add_placeholder(state, import_item.path)
			continue
		}
		target_path, target_ok := resolve_import_path(state, canonical, import_item.path)
		target := load_header(state, target_path, import_item.span) if filepath.ext(import_item.path) == ".h" else load_package(state, target_path, import_item.span)
		state.module.imports[import_id].target = target
		if !target_ok || target == ast.INVALID_PACKAGE || !state.module.packages[target].available {
			state.module.imports[import_id].valid = false
		}
		delete(target_path, state.allocator)
	}
	return pkg_id
}

declaration_conflicts :: proc(module: ^ast.Module, pkg: ast.Package_Id, name: symbol.Id) -> bool {
	for function in module.functions {
		if function.pkg == pkg && function.name == name {
			return true
		}
	}
	for global in module.globals {
		if global.pkg == pkg && global.name == name {
			return true
		}
	}
	return false
}

validate_imports :: proc(state: ^State) {
	for import_item, import_id in state.module.imports {
		if !symbol.is_valid(import_item.alias) && import_item.target != ast.INVALID_PACKAGE {
			state.module.imports[import_id].alias = state.module.packages[import_item.target].name
		}
		alias := state.module.imports[import_id].alias
		alias_text := symbol.resolve(state.symbols, alias)
		if !is_identifier(alias_text) {
			state.module.imports[import_id].diagnostic = source.add(
				state.diagnostics,
				import_item.span,
				"import requires an explicit valid identifier alias",
			)
			state.module.imports[import_id].valid = false
		}
		if declaration_conflicts(state.module, import_item.pkg, alias) {
			state.module.imports[import_id].diagnostic = source.addf(
				state.diagnostics,
				import_item.span,
				"import alias '%s' conflicts with a package declaration",
				alias_text,
			)
			state.module.imports[import_id].valid = false
		}
		for previous in state.module.imports[:import_id] {
			if previous.file == import_item.file && previous.alias == alias {
				state.module.imports[import_id].diagnostic = source.addf(
					state.diagnostics,
					import_item.span,
					"duplicate import alias '%s' in the same file",
					alias_text,
				)
				state.module.imports[import_id].valid = false
				break
			}
		}
	}
}

find_type_import :: proc(module: ^ast.Module, file: ast.File_Id, alias: symbol.Id) -> ast.Import_Id {
	for import_item, index in module.imports {
		if import_item.file == file && import_item.alias == alias {
			return ast.import_id(index)
		}
	}
	return ast.INVALID_IMPORT
}

canonical_type :: proc(
	module: ^ast.Module,
	value: types.Type,
	mapping: []types.Type,
	visiting: []bool,
) -> types.Type {
	if value < types.DYNAMIC_START {
		return value
	}
	index := int(value-types.DYNAMIC_START)
	if index < 0 || index >= len(mapping) {
		return value
	}
	if types.is_valid(mapping[index]) {
		return mapping[index]
	}
	if visiting[index] {
		return value
	}
	visiting[index] = true
	defer visiting[index] = false
	item := module.type_store.nodes[index]
	if item.kind == .Named {
		if item.qualifier != 0 {
			import_id := find_type_import(module, ast.File_Id(item.file), symbol.Id(item.qualifier))
			if import_id != ast.INVALID_IMPORT {
				module.imports[import_id].used = true
				import_item := module.imports[import_id]
				resolved := types.find_named(&module.type_store, u32(import_item.target), item.name)
				if types.is_valid(resolved) {
					mapping[index] = canonical_type(module, resolved, mapping, visiting)
					return mapping[index]
				}
			}
		}
		mapping[index] = value
		return value
	}
	if item.kind == .Alias {
		resolved := canonical_type(module, item.child, mapping, visiting)
		mapping[index] = value if !types.is_valid(resolved) else resolved
		return mapping[index]
	}
	if item.kind == .Struct || item.kind == .Union {
		mapping[index] = value
		fields := types.fields_for(&module.type_store, value)
		for &field in fields {
			field.type = canonical_type(module, field.type, mapping, visiting)
		}
		return value
	}
	if item.kind == .Function {
		params := make([]types.Type, int(item.field_count), context.temp_allocator)
		for param, param_index in types.params_for(&module.type_store, value) {
			params[param_index] = canonical_type(module, param.type, mapping, visiting)
		}
		result := canonical_type(module, item.child, mapping, visiting)
		resolved := types.function(&module.type_store, params, result, item.c_abi, item.variadic)
		mapping[index] = resolved
		return resolved
	}
	if types.is_valid(item.child) {
		item.child = canonical_type(module, item.child, mapping, visiting)
	}
	resolved := types.intern(&module.type_store, item)
	mapping[index] = resolved
	return resolved
}

canonicalize_types :: proc(module: ^ast.Module, allocator: mem.Allocator) {
	original_count := len(module.type_store.nodes)
	mapping := make([]types.Type, original_count, allocator)
	visiting := make([]bool, original_count, allocator)
	defer delete(mapping, allocator)
	defer delete(visiting, allocator)
	for &function in module.functions {
		for &param in function.params {
			param.type = canonical_type(module, param.type, mapping, visiting)
		}
		function.result = canonical_type(module, function.result, mapping, visiting)
	}
	for &global in module.globals {
		global.type = canonical_type(module, global.type, mapping, visiting)
	}
	for &statement in module.statements {
		statement.type = canonical_type(module, statement.type, mapping, visiting)
	}
	for index := 0; index < original_count; index += 1 {
		_ = canonical_type(module, types.DYNAMIC_START+types.Type(index), mapping, visiting)
	}
}

load :: proc(
	root_path: string,
	sources: ^source.Store,
	diagnostics: ^source.Diagnostics,
	symbols: ^symbol.Table,
	token_allocator := context.allocator,
	allocator := context.allocator,
	c_options := cimport.Options{},
	selected := target.DEFAULT,
) -> (ast.Module, bool) {
	module := ast.init_module(allocator)
	state := State{
		module=&module,
		sources=sources,
		diagnostics=diagnostics,
		symbols=symbols,
		token_allocator=token_allocator,
		allocator=allocator,
		c_options=c_options,
		selected=selected,
	}
	state.record_identities.allocator = allocator
	state.record_types.allocator = allocator
	defer {
		for identity in state.record_identities {
			delete(identity, allocator)
		}
		delete(state.record_identities)
		delete(state.record_types)
	}
	root := load_package(&state, root_path, source.Span{}, true)
	if root != ast.Package_Id(0) && root != ast.INVALID_PACKAGE {
		state.root_failed = true
	}
	validate_imports(&state)
	canonicalize_types(&module, allocator)
	return module, !state.root_failed
}