package source

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

Source_Id :: distinct u32
Diagnostic_Id :: distinct u32
Offset :: distinct u32

INVALID_SOURCE :: Source_Id(0xffff_ffff)
INVALID_DIAGNOSTIC :: Diagnostic_Id(0xffff_ffff)

source_id :: proc(index: int) -> Source_Id {
	assert(index >= 0 && u64(index) < u64(INVALID_SOURCE))
	return Source_Id(index)
}

diagnostic_id :: proc(index: int) -> Diagnostic_Id {
	assert(index >= 0 && u64(index) < u64(INVALID_DIAGNOSTIC))
	return Diagnostic_Id(index)
}

source_index :: proc(id: Source_Id, count: int) -> (int, bool) {
	index := int(id)
	return index, id != INVALID_SOURCE && index < count
}

diagnostic_index :: proc(id: Diagnostic_Id, count: int) -> (int, bool) {
	index := int(id)
	return index, id != INVALID_DIAGNOSTIC && index < count
}

fits_source_length :: proc(length: u64) -> bool {
	return length <= u64(max(Offset))
}

Span :: struct {
	file:  Source_Id,
	start: Offset,
	end:   Offset,
}

Source :: struct {
	id:          Source_Id,
	path:        string,
	text:        string,
	line_starts: []Offset,
}

Store :: struct {
	items:     [dynamic]Source,
	allocator: mem.Allocator,
}

Diagnostic :: struct {
	span:     Span,
	message:  string,
	severity: Severity,
}

Annotation_Kind :: enum u8 {
	Primary,
	Secondary,
	Note,
	Help,
}

Annotation :: struct {
	owner:   Diagnostic_Id,
	span:    Span,
	message: string,
	kind:    Annotation_Kind,
}

Severity :: enum u8 {
	Error,
	Warning,
}

Diagnostics :: struct {
	source:    ^Source,
	store:     ^Store,
	items:     [dynamic]Diagnostic,
	annotations: [dynamic]Annotation,
	index:     map[Diagnostic_Key]Diagnostic_Id,
	allocator: mem.Allocator,
}

Diagnostic_Key :: struct {
	span:     Span,
	message:  string,
	severity: Severity,
}

init_store :: proc(allocator := context.allocator) -> Store {
	result: Store
	result.allocator = allocator
	result.items.allocator = allocator
	return result
}

destroy_store :: proc(store: ^Store) {
	for item in store.items {
		delete(item.path, store.allocator)
		delete(item.text, store.allocator)
		delete(item.line_starts, store.allocator)
	}
	delete(store.items)
}

make_line_starts :: proc(text: string, allocator: mem.Allocator) -> []Offset {
	assert(fits_source_length(u64(len(text))))
	result: [dynamic]Offset
	result.allocator = allocator
	append(&result, 0)
	for value, offset in transmute([]byte)text {
		if value == '\n' {
			append(&result, Offset(offset+1))
		}
	}
	return result[:]
}

add_source_owned :: proc(store: ^Store, path: string, text: []byte) -> Source_Id {
	assert(fits_source_length(u64(len(text))))
	id := source_id(len(store.items))
	owned_text := string(text)
	append(&store.items, Source{
		id=id,
		path=fmt.aprintf("%s", path, allocator=store.allocator),
		text=owned_text,
		line_starts=make_line_starts(owned_text, store.allocator),
	})
	return id
}

add_source :: proc(store: ^Store, path, text: string) -> Source_Id {
	owned := make([]byte, len(text), store.allocator)
	copy(owned, transmute([]byte)text)
	return add_source_owned(store, path, owned)
}

init_diagnostics :: proc(source_file: ^Source, allocator := context.allocator) -> Diagnostics {
	result: Diagnostics
	result.source = source_file
	result.allocator = allocator
	result.items.allocator = allocator
	result.annotations.allocator = allocator
	result.index.allocator = allocator
	return result
}

init_store_diagnostics :: proc(store: ^Store, allocator := context.allocator) -> Diagnostics {
	result: Diagnostics
	result.store = store
	result.allocator = allocator
	result.items.allocator = allocator
	result.annotations.allocator = allocator
	result.index.allocator = allocator
	return result
}

destroy_diagnostics :: proc(diagnostics: ^Diagnostics) {
	delete(diagnostics.index)
	for diagnostic in diagnostics.items {
		delete(diagnostic.message, diagnostics.allocator)
	}
	for annotation in diagnostics.annotations {
		delete(annotation.message, diagnostics.allocator)
	}
	delete(diagnostics.items)
	delete(diagnostics.annotations)
}

add_with_severity :: proc(diagnostics: ^Diagnostics, span: Span, message: string, severity: Severity) -> Diagnostic_Id {
	key := Diagnostic_Key{span=span, message=message, severity=severity}
	if id, ok := diagnostics.index[key]; ok {
		return id
	}
	id := diagnostic_id(len(diagnostics.items))
	cloned := fmt.aprintf("%s", message, allocator=diagnostics.allocator)
	append(&diagnostics.items, Diagnostic{span=span, message=cloned, severity=severity})
	diagnostics.index[Diagnostic_Key{span=span, message=cloned, severity=severity}] = id
	return id
}

add :: proc(diagnostics: ^Diagnostics, span: Span, message: string) -> Diagnostic_Id {
	return add_with_severity(diagnostics, span, message, .Error)
}

add_warning :: proc(diagnostics: ^Diagnostics, span: Span, message: string) -> Diagnostic_Id {
	return add_with_severity(diagnostics, span, message, .Warning)
}

addf_with_severity :: proc(diagnostics: ^Diagnostics, span: Span, severity: Severity, format: string, args: ..any) -> Diagnostic_Id {
	message := fmt.aprintf(format, ..args, allocator=diagnostics.allocator)
	key := Diagnostic_Key{span=span, message=message, severity=severity}
	if id, ok := diagnostics.index[key]; ok {
		delete(message, diagnostics.allocator)
		return id
	}
	id := diagnostic_id(len(diagnostics.items))
	append(&diagnostics.items, Diagnostic{span=span, message=message, severity=severity})
	diagnostics.index[Diagnostic_Key{span=span, message=message, severity=severity}] = id
	return id
}

addf :: proc(diagnostics: ^Diagnostics, span: Span, format: string, args: ..any) -> Diagnostic_Id {
	return addf_with_severity(diagnostics, span, .Error, format, ..args)
}

addf_warning :: proc(diagnostics: ^Diagnostics, span: Span, format: string, args: ..any) -> Diagnostic_Id {
	return addf_with_severity(diagnostics, span, .Warning, format, ..args)
}

add_annotation :: proc(
	diagnostics: ^Diagnostics,
	owner: Diagnostic_Id,
	kind: Annotation_Kind,
	span: Span,
	message: string,
) {
	if _, ok := diagnostic_index(owner, len(diagnostics.items)); !ok {
		return
	}
	for annotation in diagnostics.annotations {
		if annotation.owner == owner && annotation.kind == kind && annotation.span == span &&
		   annotation.message == message {
			return
		}
	}
	append(&diagnostics.annotations, Annotation{
		owner=owner,
		kind=kind,
		span=span,
		message=fmt.aprintf("%s", message, allocator=diagnostics.allocator),
	})
}

set_primary_label :: proc(diagnostics: ^Diagnostics, owner: Diagnostic_Id, message: string) {
	index, ok := diagnostic_index(owner, len(diagnostics.items))
	if !ok {
		return
	}
	add_annotation(diagnostics, owner, .Primary, diagnostics.items[index].span, message)
}

add_secondary_label :: proc(diagnostics: ^Diagnostics, owner: Diagnostic_Id, span: Span, message: string) {
	add_annotation(diagnostics, owner, .Secondary, span, message)
}

add_note :: proc(diagnostics: ^Diagnostics, owner: Diagnostic_Id, message: string) {
	add_annotation(diagnostics, owner, .Note, {}, message)
}

add_help :: proc(diagnostics: ^Diagnostics, owner: Diagnostic_Id, message: string) {
	add_annotation(diagnostics, owner, .Help, {}, message)
}

line_and_column :: proc(source_file: ^Source, offset: Offset) -> (line, column: int) {
	if len(source_file.line_starts) > 0 {
		limit := min(int(offset), len(source_file.text))
		low := 0
		high := len(source_file.line_starts)
		for low < high {
			middle := low + (high-low)/2
			if int(source_file.line_starts[middle]) <= limit {
				low = middle+1
			} else {
				high = middle
			}
		}
		line = max(low, 1)
		column = limit-int(source_file.line_starts[line-1])+1
		return
	}
	line = 1
	column = 1
	limit := min(int(offset), len(source_file.text))
	for byte_value in transmute([]byte)source_file.text[:limit] {
		if byte_value == '\n' {
			line += 1
			column = 1
		} else {
			column += 1
		}
	}
	return
}

source_for_span :: proc(diagnostics: ^Diagnostics, span: Span) -> ^Source {
	if diagnostics.store != nil {
		if index, ok := source_index(span.file, len(diagnostics.store.items)); ok {
			return &diagnostics.store.items[index]
		}
	}
	return diagnostics.source
}

annotation_for :: proc(diagnostics: ^Diagnostics, owner: Diagnostic_Id, kind: Annotation_Kind) -> (Annotation, bool) {
	for annotation in diagnostics.annotations {
		if annotation.owner == owner && annotation.kind == kind {
			return annotation, true
		}
	}
	return {}, false
}

line_bounds :: proc(source_file: ^Source, line: int) -> (start, end: int, ok: bool) {
	if line <= 0 {
		return 0, 0, false
	}
	if len(source_file.line_starts) == 0 {
		current := 1
		start = 0
		for value, index in transmute([]byte)source_file.text {
			if current == line && value == '\n' {
				end = index
				if end > start && source_file.text[end-1] == '\r' {
					end -= 1
				}
				return start, end, true
			}
			if value == '\n' {
				current += 1
				start = index+1
			}
		}
		if current == line {
			return start, len(source_file.text), true
		}
		return 0, 0, false
	}
	if line > len(source_file.line_starts) {
		return 0, 0, false
	}
	start = int(source_file.line_starts[line-1])
	end = len(source_file.text)
	if line < len(source_file.line_starts) {
		end = int(source_file.line_starts[line])-1
	}
	if end > start && source_file.text[end-1] == '\r' {
		end -= 1
	}
	return start, end, true
}

write_spaces :: proc(builder: ^strings.Builder, count: int) {
	for _ in 0..<max(count, 0) {
		strings.write_byte(builder, ' ')
	}
}

write_expanded :: proc(builder: ^strings.Builder, text: string, start_column := 0) -> int {
	column := start_column
	for value in transmute([]byte)text {
		if value == '\t' {
			width := 4-column%4
			write_spaces(builder, width)
			column += width
		} else {
			strings.write_byte(builder, value)
			column += 1
		}
	}
	return column
}

display_width :: proc(text: string, start_column := 0) -> int {
	column := start_column
	for value in transmute([]byte)text {
		column += 4-column%4 if value == '\t' else 1
	}
	return column-start_column
}

decimal_width :: proc(value: int) -> int {
	width := 1
	for remaining := value; remaining >= 10; remaining /= 10 {
		width += 1
	}
	return width
}

write_excerpt :: proc(
	builder: ^strings.Builder,
	diagnostics: ^Diagnostics,
	span: Span,
	label: string,
	primary: bool,
) -> bool {
	if span == (Span{}) {
		return false
	}
	source_file := source_for_span(diagnostics, span)
	if source_file == nil {
		return false
	}
	line, column := line_and_column(source_file, span.start)
	line_start, line_end, ok := line_bounds(source_file, line)
	if !ok {
		return false
	}
	prefix := " -->" if primary else " :::"
	fmt.sbprintf(builder, "%s %s:%d:%d\n", prefix, source_file.path, line, column)
	gutter := decimal_width(line)
	write_spaces(builder, gutter+1)
	strings.write_string(builder, "|\n")
	fmt.sbprintf(builder, "%*d | ", gutter, line)
	_ = write_expanded(builder, source_file.text[line_start:line_end])
	strings.write_byte(builder, '\n')
	write_spaces(builder, gutter+1)
	strings.write_string(builder, "| ")
	start := clamp(int(span.start), line_start, line_end)
	indent := display_width(source_file.text[line_start:start])
	width := 1
	if start < line_end {
		end := clamp(int(span.end), start+1, line_end)
		width = max(display_width(source_file.text[start:end], indent), 1)
	}
	write_spaces(builder, indent)
	marker := u8('^') if primary else u8('-')
	for _ in 0..<width {
		strings.write_byte(builder, marker)
	}
	if len(label) > 0 {
		strings.write_byte(builder, ' ')
		strings.write_string(builder, label)
	}
	strings.write_byte(builder, '\n')
	return true
}

format :: proc(diagnostics: ^Diagnostics, id: Diagnostic_Id, allocator := context.allocator) -> string {
	index, ok := diagnostic_index(id, len(diagnostics.items))
	if !ok {
		path := "<unknown>"
		if diagnostics.source != nil {
			path = diagnostics.source.path
		} else if diagnostics.store != nil && len(diagnostics.store.items) > 0 {
			path = diagnostics.store.items[0].path
		}
		return fmt.aprintf("%s: compiler recovery error", path, allocator=allocator)
	}
	diagnostic := diagnostics.items[index]
	source_file := source_for_span(diagnostics, diagnostic.span)
	severity := "warning" if diagnostic.severity == .Warning else "error"
	if diagnostic.span == (Span{}) {
		path := source_file.path if source_file != nil else "<unknown>"
		return fmt.aprintf("%s: %s: %s", path, severity, diagnostic.message, allocator=allocator)
	}
	if source_file == nil {
		return fmt.aprintf("<unknown>: %s: %s", severity, diagnostic.message, allocator=allocator)
	}
	builder := strings.builder_make(allocator)
	fmt.sbprintf(&builder, "%s: %s\n", severity, diagnostic.message)
	primary_label := ""
	if annotation, found := annotation_for(diagnostics, id, .Primary); found {
		primary_label = annotation.message
	}
	_ = write_excerpt(&builder, diagnostics, diagnostic.span, primary_label, true)
	for annotation in diagnostics.annotations {
		if annotation.owner == id && annotation.kind == .Secondary {
			_ = write_excerpt(&builder, diagnostics, annotation.span, annotation.message, false)
		}
	}
	for annotation in diagnostics.annotations {
		if annotation.owner != id {
			continue
		}
		#partial switch annotation.kind {
		case .Note:
			fmt.sbprintf(&builder, "note: %s\n", annotation.message)
		case .Help:
			fmt.sbprintf(&builder, "help: %s\n", annotation.message)
		case:
		}
	}
	result := strings.to_string(builder)
	if len(result) > 0 && result[len(result)-1] == '\n' {
		return result[:len(result)-1]
	}
	return result
}

print_all :: proc(diagnostics: ^Diagnostics) {
	for _, id in diagnostics.items {
		message := format(diagnostics, diagnostic_id(id))
		fmt.eprintln(message)
		delete(message)
	}
}