package main

import (
	"html"
	"html/template"
	"path"
	"strings"

	"github.com/alecthomas/chroma/v2"
	"github.com/alecthomas/chroma/v2/lexers"
)

var radianceLexer = chroma.MustNewLexer(
	&chroma.Config{Name: "Radiance"},
	func() chroma.Rules {
		return chroma.Rules{"root": {
			{Pattern: `//[^\n]*`, Type: chroma.CommentSingle},
			{Pattern: `"(?:\\.|[^"\\])*"|'[^']{1,2}'`, Type: chroma.LiteralString},
			{Pattern: `\b(?:0x[0-9a-fA-F]+|0b[01]+|[0-9]+(?:\.[0-9]+)?)\b`, Type: chroma.LiteralNumber},
			{Pattern: `@[a-zA-Z]+`, Type: chroma.NameBuiltin},
			{Pattern: `\b(?:fn|if|else|for|while|break|switch|match|set|record|union|constant|align|let|use|mod|module|case|continue|return|true|false|loop|extern|panic|device|register|bit|catch|throw|throws|test|trait|instance|at|mut|nil|undefined|static|in|is|where|export|as|and|or|xor|not|try|atomic|select|of|assert)\b`, Type: chroma.Keyword},
			{Pattern: `\b(?:bool|u8|u16|u32|u64|i8|i16|i32|i64|f32|void|opaque)\b`, Type: chroma.KeywordType},
			{Pattern: `->|=>|::|\.\.|<>|\?{1,2}|!{1,2}|>=?|<=?|[&*.=+\-/%]`, Type: chroma.Operator},
			{Pattern: `[(){}\[\]]`, Type: chroma.Punctuation},
			{Pattern: `\s+`, Type: chroma.TextWhitespace},
			{Pattern: `.`, Type: chroma.Text},
		}}
	},
)

var radianceILLexer = chroma.MustNewLexer(
	&chroma.Config{Name: "Radiance IL"},
	func() chroma.Rules {
		return chroma.Rules{"root": {
			{Pattern: `//[^\n]*`, Type: chroma.CommentSingle},
			{Pattern: `"[^"\n]*"`, Type: chroma.LiteralString},
			{Pattern: `%[0-9]+`, Type: chroma.NameVariable},
			{Pattern: `@[A-Za-z_][A-Za-z#0-9_]*`, Type: chroma.NameLabel},
			{Pattern: `\$[A-Za-z_][A-Za-z0-9_]*`, Type: chroma.Name},
			{Pattern: `\b(?:-?[0-9]+|0x[a-fA-F0-9]+)\b`, Type: chroma.LiteralNumber},
			{Pattern: `\b(?:fn|data|extern|mut|align)\b`, Type: chroma.Keyword},
			{Pattern: `\b(?:br\.eq|br\.ne|br\.slt|br\.ult|reserve|load|sload|store|blit|copy|add|sub|mul|sdiv|udiv|srem|urem|neg|eq|ne|slt|sge|ult|uge|and|or|xor|shl|sshr|ushr|not|zext|sext|call|ret|jmp|switch|unreachable|ecall|ebreak)\b`, Type: chroma.NameBuiltin},
			{Pattern: `\b(?:str|sym|undef|w8|w16|w32|w64)\b`, Type: chroma.KeywordType},
			{Pattern: `[{}();:,]`, Type: chroma.Punctuation},
			{Pattern: `\s+`, Type: chroma.TextWhitespace},
			{Pattern: `.`, Type: chroma.Text},
		}}
	},
)

// highlightLines returns one escaped HTML fragment for each source line.
func highlightLines(filename string, lines []string) []template.HTML {
	if len(lines) == 0 {
		return nil
	}

	lexer := lexerForFilename(filename)
	iterator, err := lexer.Tokenise(nil, strings.Join(lines, "\n"))
	if err != nil {
		return escapeLines(lines)
	}

	highlighted := make([]strings.Builder, len(lines))
	line := 0
	for token := iterator(); token != chroma.EOF; token = iterator() {
		parts := strings.Split(token.Value, "\n")
		class := tokenClass(token.Type)
		for i, part := range parts {
			if i > 0 {
				line++
			}
			if line >= len(highlighted) || part == "" {
				continue
			}
			if class != "" {
				highlighted[line].WriteString(`<span class="`)
				highlighted[line].WriteString(class)
				highlighted[line].WriteString(`">`)
			}
			highlighted[line].WriteString(html.EscapeString(part))
			if class != "" {
				highlighted[line].WriteString(`</span>`)
			}
		}
	}

	result := make([]template.HTML, len(highlighted))
	for i := range highlighted {
		result[i] = template.HTML(highlighted[i].String())
	}
	return result
}

func highlightDiff(file DiffFile, lines []DiffLine) {
	oldSource := make([]string, 0, len(lines))
	newSource := make([]string, 0, len(lines))
	for i := range lines {
		if lines[i].OldNum > 0 {
			oldSource = append(oldSource, lines[i].Content)
		}
		if lines[i].NewNum > 0 {
			newSource = append(newSource, lines[i].Content)
		}
	}

	oldHighlighted := highlightLines(file.OldName, oldSource)
	newHighlighted := highlightLines(file.NewName, newSource)
	oldIndex, newIndex := 0, 0
	for i := range lines {
		line := &lines[i]
		if line.Type == "del" {
			line.Highlighted = oldHighlighted[oldIndex]
		} else {
			line.Highlighted = newHighlighted[newIndex]
		}
		if line.OldNum > 0 {
			oldIndex++
		}
		if line.NewNum > 0 {
			newIndex++
		}
	}
}

func lexerForFilename(filename string) chroma.Lexer {
	switch strings.ToLower(path.Ext(filename)) {
	case ".rad":
		return radianceLexer
	case ".ril":
		return radianceILLexer
	}
	if lexer := lexers.Match(filename); lexer != nil {
		return chroma.Coalesce(lexer)
	}
	return lexers.Fallback
}

func tokenClass(tokenType chroma.TokenType) string {
	switch {
	case tokenType.InCategory(chroma.Keyword):
		return "syntax-keyword"
	case tokenType.InSubCategory(chroma.NameBuiltin):
		return "syntax-builtin"
	case tokenType.InSubCategory(chroma.LiteralString):
		return "syntax-string"
	case tokenType.InSubCategory(chroma.LiteralNumber):
		return "syntax-number"
	case tokenType.InCategory(chroma.Comment):
		return "syntax-comment"
	case tokenType.InCategory(chroma.Operator):
		return "syntax-operator"
	case tokenType.InCategory(chroma.Punctuation):
		return "syntax-punctuation"
	case tokenType.InSubCategory(chroma.NameVariable), tokenType == chroma.NameLabel:
		return "syntax-name"
	default:
		return ""
	}
}

func escapeLines(lines []string) []template.HTML {
	escaped := make([]template.HTML, len(lines))
	for i, line := range lines {
		escaped[i] = template.HTML(html.EscapeString(line))
	}
	return escaped
}