package docsite

import (
	"html"
	"slices"
	"strings"

	"github.com/alecthomas/chroma/v2"
	"github.com/alecthomas/chroma/v2/lexers"
	"github.com/yuin/goldmark/ast"
	"github.com/yuin/goldmark/renderer"
	"github.com/yuin/goldmark/util"
)

// codeRenderer renders every fenced code block as the UI-SPEC code block:
// <figure class="code">, a <figcaption> naming the source of a src= fence
// (linked to source_url), a copy button that stays hidden until site.js
// finds a clipboard, and <pre><code> highlighted at build time. Tokens
// come from chroma/v2 lexers and are mapped onto the tok-* classes; the
// chroma HTML formatter and its styles are not used, so the output has no
// inline style and only the UI-SPEC class names.
type codeRenderer struct{}

func (codeRenderer) RegisterFuncs(r renderer.NodeRendererFuncRegisterer) {
	r.Register(ast.KindFencedCodeBlock, renderFence)
}

func attrString(n ast.Node, name string) string {
	v, ok := n.AttributeString(name)
	if !ok {
		return ""
	}
	b, _ := v.([]byte)
	return string(b)
}

func renderFence(w util.BufWriter, src []byte, node ast.Node, entering bool) (ast.WalkStatus, error) {
	if !entering {
		return ast.WalkContinue, nil
	}
	n := node.(*ast.FencedCodeBlock)
	_, _ = w.WriteString(`<figure class="code">`)
	if ref := attrString(n, "data-src"); ref != "" {
		_, _ = w.WriteString("<figcaption>")
		if href := attrString(n, "data-href"); href != "" {
			_, _ = w.WriteString(`<a href="` + html.EscapeString(href) + `">` + html.EscapeString(ref) + "</a>")
		} else {
			_, _ = w.WriteString(html.EscapeString(ref))
		}
		_, _ = w.WriteString("</figcaption>")
	}
	_, _ = w.WriteString(`<button type="button" class="copy-button" aria-label="Copy code" hidden>`)
	writeIcon(w, "copy")
	writeIcon(w, "check")
	_, _ = w.WriteString("</button><pre><code")
	lang := ""
	if l := n.Language(src); l != nil {
		lang = string(l)
		_, _ = w.WriteString(` class="language-` + html.EscapeString(lang) + `"`)
	}
	_ = w.WriteByte('>')
	var code strings.Builder
	for i := 0; i < n.Lines().Len(); i++ {
		line := n.Lines().At(i)
		code.Write(line.Value(src))
	}
	_, _ = w.WriteString(highlight(lang, code.String()))
	_, _ = w.WriteString("</code></pre></figure>\n")
	return ast.WalkSkipChildren, nil
}

// highlight returns code as escaped HTML with tok-* spans. Unknown
// languages, plain text and lexer errors fall back to escaped text.
func highlight(lang, code string) string {
	var b strings.Builder
	if slices.Contains(shellLangs, lang) {
		highlightShell(&b, code)
		return b.String()
	}
	lexer := lexerFor(lang)
	if lexer == nil || !writeTokens(&b, lexer, lang, code) {
		return html.EscapeString(code)
	}
	return b.String()
}

// lexerFor returns the chroma lexer for a fence language, or nil for
// plain text and unknown languages.
func lexerFor(lang string) chroma.Lexer {
	switch lang {
	case "", "text", "txt", "plain", "plaintext":
		return nil
	}
	l := lexers.Get(lang)
	if l == nil {
		return nil
	}
	return chroma.Coalesce(l)
}

// writeTokens tokenises code and writes each token as a tok-* span or as
// escaped plain text. It reports false when the lexer fails.
func writeTokens(b *strings.Builder, lexer chroma.Lexer, lang, code string) bool {
	it, err := lexer.Tokenise(nil, code)
	if err != nil {
		return false
	}
	tokens := it.Tokens()
	// Lexers configured with EnsureNL append a newline the source did not
	// have; drop it so a shell line fragment stays one line.
	if n := len(tokens); n > 0 && !strings.HasSuffix(code, "\n") {
		tokens[n-1].Value = strings.TrimSuffix(tokens[n-1].Value, "\n")
	}
	for _, tok := range tokens {
		writeSpan(b, tokClass(tok.Type, lang), tok.Value)
	}
	return true
}

func writeSpan(b *strings.Builder, class, value string) {
	if value == "" {
		return
	}
	if class == "" {
		b.WriteString(html.EscapeString(value))
		return
	}
	b.WriteString(`<span class="` + class + `">` + html.EscapeString(value) + "</span>")
}

// tokClass maps a chroma token type onto a UI-SPEC syntax class, or "" for
// plain text (identifiers, operators, punctuation).
func tokClass(t chroma.TokenType, lang string) string {
	data := lang == "yaml" || lang == "yml" || lang == "json"
	switch {
	case t == chroma.GenericPrompt:
		return "tok-prompt"
	case t.InCategory(chroma.Comment):
		return "tok-com"
	case t == chroma.NameTag || t == chroma.NameAttribute:
		return "tok-key"
	case t == chroma.KeywordConstant && data:
		// YAML and JSON booleans and null.
		return "tok-num"
	case t.InCategory(chroma.Keyword):
		return "tok-kw"
	case t.InSubCategory(chroma.LiteralString):
		return "tok-str"
	case t.InSubCategory(chroma.LiteralNumber):
		return "tok-num"
	case t == chroma.Literal && data:
		// A plain YAML scalar is a string.
		return "tok-str"
	}
	return ""
}

// highlightShell highlights a shell session line by line: a leading "$ "
// is a tok-prompt (never copied), a "#" line is a comment, the command
// word is tok-kw and the rest goes through the bash lexer. Continuation
// lines (after a trailing "\") have no command word.
func highlightShell(b *strings.Builder, code string) {
	lexer := lexerFor("bash")
	cont := false
	lines := strings.SplitAfter(code, "\n")
	for _, line := range lines {
		body := strings.TrimSuffix(line, "\n")
		nl := len(body) < len(line)
		rest := body
		if !cont {
			trimmed := strings.TrimLeft(body, " \t")
			b.WriteString(html.EscapeString(body[:len(body)-len(trimmed)]))
			rest = trimmed
			if p, ok := strings.CutPrefix(rest, "$ "); ok {
				writeSpan(b, "tok-prompt", "$ ")
				rest = p
			}
			if strings.HasPrefix(rest, "#") {
				writeSpan(b, "tok-com", rest)
				rest = ""
			} else if rest != "" {
				word, tail, found := strings.Cut(rest, " ")
				writeSpan(b, "tok-kw", word)
				if found {
					b.WriteString(" ")
				}
				rest = tail
			}
		}
		if rest != "" {
			if lexer == nil || !writeTokens(b, lexer, "bash", rest) {
				b.WriteString(html.EscapeString(rest))
			}
		}
		cont = strings.HasSuffix(strings.TrimRight(body, " \t"), `\`)
		if nl {
			b.WriteString("\n")
		}
	}
}