// Package highlight is the single place that turns source code into HTML.
// File pages, markdown fenced blocks and README rendering all call HTML, so
// identical input always yields identical markup regardless of caller.
package highlight

import (
	"bytes"
	"embed"
	"errors"
	"fmt"
	"html"
	"html/template"
	"io"
	"io/fs"
	"maps"
	"os"
	"path/filepath"
	"reflect"
	"slices"
	"strings"
	"time"

	"github.com/alecthomas/chroma/v2"
	chromahtml "github.com/alecthomas/chroma/v2/formatters/html"
	"github.com/alecthomas/chroma/v2/lexers"
	"github.com/alecthomas/chroma/v2/styles"
)

// Custom describes a lexer configured in repo-render.yaml.
type Custom struct {
	Name    string
	Globs   []string
	Aliases []string
}

type globEntry struct {
	glob  string
	lexer chroma.Lexer
}

// Registry resolves lexers by path or name and renders code.
type Registry struct {
	style   *chroma.Style
	light   *chroma.Style
	globs   []globEntry
	aliases map[string]chroma.Lexer
	plain   chroma.Lexer

	// markdown is chroma's markdown lexer, but looking up fenced-block
	// languages in fences rather than chroma's global registry.
	markdown chroma.Lexer
	fences   *chroma.LexerRegistry
}

// New builds a registry. Custom lexers are looked up as <dir>/<name>.xml in
// each of lexerDirs in order, then as chroma builtins; a name matching
// neither is an error. Missing directories are skipped.
func New(lexerDirs []string, custom []Custom, styleName string) (*Registry, error) {
	st := styles.Fallback
	if styleName != "" {
		var err error
		if st, err = findStyle(styleName); err != nil {
			return nil, err
		}
	}
	r := &Registry{
		style:   st,
		aliases: map[string]chroma.Lexer{},
		plain:   lexers.Get("plaintext"),
	}
	r.fences = chroma.NewLexerRegistry()
	md, err := fenceMarkdown(r.fences, markdownInline()...)
	if err != nil {
		return nil, err
	}
	r.markdown = md
	// Globs are tried in name order whatever order the caller has, and
	// the caller's slice is left as it is.
	custom = slices.SortedFunc(slices.Values(custom), func(a, b Custom) int { return strings.Compare(a.Name, b.Name) })
	for _, c := range custom {
		lx, err := loadLexer(lexerDirs, c.Name)
		if err != nil {
			return nil, err
		}
		for _, g := range c.Globs {
			r.globs = append(r.globs, globEntry{glob: g, lexer: lx})
		}
		r.aliases[strings.ToLower(c.Name)] = lx
		for _, a := range c.Aliases {
			r.aliases[strings.ToLower(a)] = lx
		}
	}
	r.fillFences()
	return r, nil
}

// Bundled holds the lexers that ship with the renderer (lexers/*.xml),
// used for a name no lexer directory has.
//
//go:embed lexers/*.xml
var Bundled embed.FS

// loadLexer returns the lexer called name: from name.xml in the first of
// lexerDirs that has one, or else the bundled one, or else chroma's
// builtin of that name.
func loadLexer(lexerDirs []string, name string) (chroma.Lexer, error) {
	for _, dir := range lexerDirs {
		if _, err := os.Stat(filepath.Join(dir, name+".xml")); err != nil {
			continue
		}
		lx, err := chroma.NewXMLLexer(os.DirFS(dir), name+".xml")
		if err != nil {
			return nil, fmt.Errorf("lexer %s: %w", name, err)
		}
		return lx, nil
	}
	if sub, err := fs.Sub(Bundled, "lexers"); err == nil {
		if _, err := fs.Stat(sub, name+".xml"); err == nil {
			lx, err := chroma.NewXMLLexer(sub, name+".xml")
			if err != nil {
				return nil, fmt.Errorf("bundled lexer %s: %w", name, err)
			}
			return lx, nil
		}
	}
	if lx := lexers.Get(name); lx != nil {
		return lx, nil
	}
	where := "the bundled lexers"
	if len(lexerDirs) > 0 {
		where = strings.Join(lexerDirs, ", ") + " or " + where
	}
	return nil, fmt.Errorf("lexer %q: no %s.xml in %s and no chroma builtin by that name", name, name, where)
}

// proseRun lexes a run of plain markdown text in one match. chroma's inline
// state ends in `.|\n`, so prose otherwise goes a character at a time, each
// only after every other rule has failed there: three times slower. The
// run stops at each character that begins another inline rule (whitespace
// for emphasis, < for comments, \ for escapes, ` for code, @ and # for
// mentions, [ for links), so it never covers a position where one could
// match. Inline is tried only after root's rules, which need the start of
// a line or, for comments, a <.
var proseRun = chroma.Rule{Pattern: "[^\\s<\\\\`@#\\[]+", Type: chroma.Text}

// lineStart has chroma's emphasis, strong and strikethrough rules again
// for the start of a line. Theirs need a space before the marker, which a
// paragraph's first word lacks, and so does a "- " or "* " list item's,
// since chroma lexes the item's text on its own. The marked text must not
// start with a space, so "* * *" stays a rule, nor hold a backtick, so
// code inside keeps its own highlighting. Mid-line, proseRun never covers
// a line start, so these only have to come before it.
var lineStart = []chroma.Rule{
	{Pattern: "^(\\*|_)(?!\\s)((?:(?!\\1)[^`\\n])+)(\\1)((?=\\W|\\n))", Type: chroma.ByGroups(chroma.GenericEmph, chroma.GenericEmph, chroma.GenericEmph, chroma.Text)},
	{Pattern: "^((\\*\\*|__)(?!\\s)[^`\\n]+?)\\2((?=\\W|\\n))", Type: chroma.ByGroups(chroma.GenericStrong, chroma.GenericStrong, chroma.Text)},
	{Pattern: "^(~~(?!\\s)[^~`]+~~)((?=\\W|\\n))", Type: chroma.ByGroups(chroma.GenericDeleted, chroma.Text)},
}

// markdownInline is what the markdown lexer tries first in its inline
// state, in order.
func markdownInline() []chroma.Rule { return append(slices.Clone(lineStart), proseRun) }

// fenceMarkdown returns a copy of chroma's markdown lexer that resolves
// fenced-block languages in reg. The original uses chroma's global
// registry, which knows nothing of custom lexers and, for a fence that
// names no known lexer (or none at all), scans every filename glob twice.
// Rules in inline are tried first in the copy's inline state. chroma's
// lexer is a frontmatter-handling wrapper around a RegexLexer; if a chroma
// upgrade changes that shape, this fails rather than quietly losing custom
// fence languages.
func fenceMarkdown(reg *chroma.LexerRegistry, inline ...chroma.Rule) (chroma.Lexer, error) {
	unexpected := errors.New("chroma's markdown lexer is not the wrapped RegexLexer fenceMarkdown expects")
	md := reflect.ValueOf(lexers.Markdown)
	if md.Kind() != reflect.Pointer || md.Elem().Kind() != reflect.Struct {
		return nil, unexpected
	}
	inner := md.Elem().FieldByName("Lexer")
	if !inner.IsValid() {
		return nil, unexpected
	}
	base, ok := reflect.TypeAssert[*chroma.RegexLexer](inner)
	if !ok {
		return nil, unexpected
	}
	rules, err := base.Rules()
	if err != nil {
		return nil, err
	}
	if len(inline) > 0 {
		own := maps.Clone(rules)
		own["inline"] = append(slices.Clone(inline), rules["inline"]...)
		rules = own
	}
	clone := chroma.MustNewLexer(base.Config(), func() chroma.Rules { return rules })
	clone.SetRegistry(reg)
	out := reflect.New(md.Elem().Type())
	out.Elem().Set(md.Elem())
	out.Elem().FieldByName("Lexer").Set(reflect.ValueOf(clone))
	lx, ok := reflect.TypeAssert[chroma.Lexer](out)
	if !ok {
		return nil, unexpected
	}
	return lx, nil
}

// fenceAlias registers a lexer in the fence registry under extra names.
// Its SetRegistry does nothing, so registering it leaves chroma's shared
// lexers pointing at chroma's registry.
type fenceAlias struct {
	chroma.Lexer
	cfg *chroma.Config
}

func (a *fenceAlias) Config() *chroma.Config                         { return a.cfg }
func (a *fenceAlias) SetRegistry(*chroma.LexerRegistry) chroma.Lexer { return a }

// fillFences registers every name ByName knows other than through a
// wildcard glob, then empties the registry's lexer list: chroma's Get
// falls back to Match over that list, so a miss now costs nothing and
// leaves the block unhighlighted, as it always did.
func (r *Registry) fillFences() {
	custom := map[chroma.Lexer]bool{}
	for _, lx := range r.aliases {
		custom[lx] = true
	}
	groups := map[chroma.Lexer]*fenceAlias{}
	var order []chroma.Lexer
	add := func(name string, lx chroma.Lexer) {
		a, ok := groups[lx]
		if !ok {
			final := lx
			if !custom[lx] {
				final = r.builtin(lx)
			}
			a = &fenceAlias{Lexer: final, cfg: &chroma.Config{Name: fmt.Sprintf("\x00fence%d", len(order))}}
			groups[lx] = a
			order = append(order, lx)
		}
		a.cfg.Aliases = append(a.cfg.Aliases, name)
	}
	for name, lx := range r.aliases {
		add(name, lx)
	}
	for _, name := range builtins().names() {
		if _, ok := r.aliases[name]; ok {
			continue
		}
		if lx := builtins().get(name); lx != nil {
			add(name, lx)
		}
	}
	for _, lx := range order {
		r.fences.Register(groups[lx])
	}
	r.fences.Lexers = nil
}

// builtin readies a lexer from chroma's registry for use, or returns nil.
func (r *Registry) builtin(lx chroma.Lexer) chroma.Lexer {
	switch lx {
	case nil:
		return nil
	case lexers.Markdown:
		return r.markdown
	}
	return lx
}

// ByPath picks a lexer for a file path: repo-render.yaml globs, then
// chroma's filename registry, then plain text. Never nil.
func (r *Registry) ByPath(path string) chroma.Lexer {
	base := filepath.Base(path)
	for _, g := range r.globs {
		if ok, _ := filepath.Match(g.glob, base); ok {
			return g.lexer
		}
	}
	if lx := r.builtin(builtins().match(base)); lx != nil {
		return lx
	}
	return r.plain
}

// ByFile is ByPath, except that a file its name leaves as plain text is
// looked up by the interpreter on its "#!" line ("#!/bin/sh",
// "#!/usr/bin/env python3"), so extensionless scripts are still code.
func (r *Registry) ByFile(path string, data []byte) chroma.Lexer {
	lx := r.ByPath(path)
	if !r.IsPlain(lx) {
		return lx
	}
	if sb := r.byShebang(data); sb != nil {
		return sb
	}
	return lx
}

// byShebang returns the lexer named by data's "#!" line, or nil.
func (r *Registry) byShebang(data []byte) chroma.Lexer {
	if !bytes.HasPrefix(data, []byte("#!")) {
		return nil
	}
	line := data[2:]
	if i := bytes.IndexByte(line, '\n'); i >= 0 {
		line = line[:i]
	}
	if len(line) > 256 {
		return nil
	}
	fields := strings.Fields(string(line))
	if len(fields) == 0 {
		return nil
	}
	name := filepath.Base(fields[0])
	if name == "env" {
		// Skip env's options and VAR=value assignments.
		name = ""
		for _, f := range fields[1:] {
			if !strings.HasPrefix(f, "-") && !strings.Contains(f, "=") {
				name = filepath.Base(f)
				break
			}
		}
	}
	// "python3.12" → "python3" → "python".
	for _, n := range []string{name, strings.TrimRight(name, "0123456789.")} {
		if n == "" {
			continue
		}
		if lx, ok := r.aliases[strings.ToLower(n)]; ok {
			return lx
		}
		if lx := r.builtin(builtins().get(n)); lx != nil && !r.IsPlain(lx) {
			return lx
		}
	}
	return nil
}

// ByName picks a lexer for a fence info string or alias. Never nil.
func (r *Registry) ByName(name string) chroma.Lexer {
	name = strings.ToLower(strings.TrimSpace(name))
	if name == "" {
		return r.plain
	}
	if lx, ok := r.aliases[name]; ok {
		return lx
	}
	if lx := r.builtin(builtins().get(name)); lx != nil {
		return lx
	}
	return r.plain
}

// Name returns the display name of a lexer.
func Name(lx chroma.Lexer) string {
	if lx == nil {
		return "Text"
	}
	return lx.Config().Name
}

// IsPlain reports whether lx is the plain text fallback.
func (r *Registry) IsPlain(lx chroma.Lexer) bool {
	return lx == nil || lx.Config().Name == r.plain.Config().Name
}

// Limits on highlighting one input, past which HTML gives up and shows it
// plain. Some of chroma's lexers can loop forever without consuming input
// (Jungle on an unclosed '"', JSONata on a control character, chroma issue
// #1377), emitting only empty tokens; others could just be very slow.
// Either would stall the render while it holds the render lock.
const (
	maxEmptyTokens = 1000
	highlightTime  = 5 * time.Second
)

// HTML renders src as a <code class="chroma"> element with one span per
// token, classes only (no inline styles), tabs preserved. Line numbers and
// anchors are the caller's job. Input a lexer cannot finish within the
// limits above, or panics on (YANG on "/*/*"), is shown unhighlighted.
func (r *Registry) HTML(src []byte, lx chroma.Lexer) template.HTML {
	if lx == nil {
		lx = r.plain
	}
	if out, ok := r.tokensHTML(src, lx); ok {
		return out
	}
	return template.HTML(`<code class="chroma">` + html.EscapeString(string(src)) + `</code>`)
}

// tokensHTML is HTML's highlighted rendering, or false when lexing fails,
// panics or passes a limit.
func (r *Registry) tokensHTML(src []byte, lx chroma.Lexer) (out template.HTML, ok bool) {
	defer func() {
		if recover() != nil {
			out, ok = "", false
		}
	}()
	it, err := lx.Tokenise(nil, string(src))
	if err != nil {
		return "", false
	}
	var b strings.Builder
	b.Grow(len(src) * 2)
	b.WriteString(`<code class="chroma">`)
	write := func(tok chroma.Token) {
		cls := class(tok.Type)
		if cls == "" {
			b.WriteString(html.EscapeString(tok.Value))
			return
		}
		b.WriteString(`<span class="`)
		b.WriteString(cls)
		b.WriteString(`">`)
		b.WriteString(html.EscapeString(tok.Value))
		b.WriteString(`</span>`)
	}
	// Tokens are merged into runs of one type as chroma.Coalesce would,
	// but here, where the limits can see every token: Coalesce drops the
	// empty ones inside its own loop, so a lexer looping on them would
	// never return from it.
	deadline := time.Now().Add(highlightTime)
	empty := 0
	var run chroma.Token
	for tok := it(); tok != chroma.EOF; tok = it() {
		if time.Now().After(deadline) {
			return "", false
		}
		if tok.Value == "" {
			if empty++; empty > maxEmptyTokens {
				return "", false
			}
			continue
		}
		empty = 0
		if run.Value != "" && run.Type == tok.Type && len(run.Value) < 8192 {
			run.Value += tok.Value
			continue
		}
		if run.Value != "" {
			write(run)
		}
		run = tok
	}
	if run.Value != "" {
		write(run)
	}
	b.WriteString(`</code>`)
	return template.HTML(b.String()), true
}

// class mirrors chroma's HTML formatter class naming so WriteCSS output
// matches: walk up the token hierarchy to the first standard class.
func class(t chroma.TokenType) string {
	for t != 0 {
		if cls, ok := chroma.StandardTypes[t]; ok {
			return cls
		}
		t = t.Parent()
	}
	return chroma.StandardTypes[t]
}

// LineCount returns the number of lines src occupies on a file page.
func LineCount(src []byte) int {
	if len(src) == 0 {
		return 0
	}
	n := bytes.Count(src, []byte{'\n'})
	if src[len(src)-1] != '\n' {
		n++
	}
	return n
}

// findStyle looks up a chroma style by name. chroma's Get answers an
// unknown name with its fallback style, so that is an error here unless it
// was the one asked for.
func findStyle(name string) (*chroma.Style, error) {
	st := styles.Get(name)
	if st == nil || (st == styles.Fallback && st.Name != name) {
		return nil, fmt.Errorf("unknown highlight style %q", name)
	}
	return st, nil
}

// SetLight selects the style used in light mode.
func (r *Registry) SetLight(name string) error {
	st, err := findStyle(name)
	if err != nil {
		return err
	}
	r.light = st
	return nil
}

// WriteCSS writes the syntax stylesheet. With no light style it is the dark
// style as is. With one, the two are mutually exclusive: layering light
// rules over dark ones would leave any property the light style does not
// set (token backgrounds, bold) leaking through. Each style is emitted for
// the user's explicit choice (:root[data-theme]) and, when they have made
// none, for the system preference.
func (r *Registry) WriteCSS(w io.Writer) error {
	f := chromahtml.New(chromahtml.WithClasses(true), chromahtml.WithAllClasses(true))
	var dark bytes.Buffer
	if err := f.WriteCSS(&dark, r.style); err != nil {
		return err
	}
	if r.light == nil {
		_, err := w.Write(dark.Bytes())
		return err
	}
	var light bytes.Buffer
	if err := f.WriteCSS(&light, r.light); err != nil {
		return err
	}
	_, err := fmt.Fprintf(w, "/* dark */\n%s@media (prefers-color-scheme: dark) {\n%s}\n/* light */\n%s@media (prefers-color-scheme: light) {\n%s}\n",
		scopeCSS(dark.String(), `:root[data-theme="dark"] `, false),
		scopeCSS(dark.String(), `:root:not([data-theme]) `, true),
		scopeCSS(light.String(), `:root[data-theme="light"] `, false),
		scopeCSS(light.String(), `:root:not([data-theme]) `, true))
	return err
}

// scopeCSS prefixes the selector of every one-line rule chroma emits
// ("/* Name */ .sel { ... }") with prefix.
func scopeCSS(css, prefix string, indent bool) string {
	var b strings.Builder
	for line := range strings.SplitSeq(css, "\n") {
		_, rule, ok := strings.Cut(line, "*/ ")
		if !ok || !strings.Contains(line, "{") {
			continue
		}
		if indent {
			b.WriteString("  ")
		}
		b.WriteString(prefix)
		b.WriteString(rule)
		b.WriteString("\n")
	}
	return b.String()
}

// StyleName returns the configured chroma style name.
func (r *Registry) StyleName() string { return r.style.Name }

// TextColour returns the style's foreground colour as "#rrggbb", or "" if
// the style does not set one.
func (r *Registry) TextColour() string { return textColour(r.style) }

// LightTextColour is TextColour for the light style ("" when none is set).
func (r *Registry) LightTextColour() string { return textColour(r.light) }

func textColour(st *chroma.Style) string {
	if st == nil {
		return ""
	}
	if e := st.Get(chroma.Text); e.Colour.IsSet() {
		return e.Colour.String()
	}
	return ""
}

// BinaryProbe is how much of the start of a file IsBinary looks at.
const BinaryProbe = 8192

// IsBinary applies git's heuristic: a NUL in the first 8 KiB means binary.
func IsBinary(data []byte) bool {
	if len(data) > BinaryProbe {
		data = data[:BinaryProbe]
	}
	return bytes.IndexByte(data, 0) >= 0
}