package highlight

import (
	"cmp"
	"path/filepath"
	"slices"
	"sort"
	"strings"
	"sync"

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

// chroma's LexerRegistry.Match tries every filename glob of every lexer,
// each again with every backup suffix: ~15,000 filepath.Match calls, about
// 8ms per lookup, and Get falls back to two of those on any unknown name.
// builtinIndex gives the same answers from hash tables: nearly all globs
// are "*.ext" or a literal name, leaving a dozen to match one by one.

// ignoredSuffixes mirrors chroma's list of backup suffixes a filename may
// carry and still match a glob ("foo.py.bak" is Python).
var ignoredSuffixes = [...]string{
	"~", ".bak", ".old", ".orig",
	".dpkg-dist", ".dpkg-old", ".ucf-dist", ".ucf-new", ".ucf-old",
	".rpmnew", ".rpmorig", ".rpmsave",
	".in",
}

// globHit is one glob that matched: its lexer's position in the registry
// and its position in that lexer's list. chroma breaks priority ties by the
// order it finds matches, so hits are replayed in this order.
type globHit struct{ lexer, glob int }

type globRule struct {
	pattern string
	hit     globHit
}

// globIndex holds one tier of globs (Filenames or AliasFilenames).
type globIndex struct {
	exact map[string][]globHit // literal names: "Makefile"
	ext   map[string][]globHit // "*.ext" globs, keyed by ".ext"
	other []globRule           // everything else: "*.[1-9]", "Makefile.*"
}

func (g *globIndex) add(pattern string, hit globHit) {
	switch {
	case !strings.ContainsAny(pattern, `*?[\`):
		g.exact[pattern] = append(g.exact[pattern], hit)
	case strings.HasPrefix(pattern, "*.") && !strings.ContainsAny(pattern[1:], `*?[\`):
		g.ext[pattern[1:]] = append(g.ext[pattern[1:]], hit)
	default:
		g.other = append(g.other, globRule{pattern, hit})
	}
}

// hits returns every glob matching name, directly or once a backup suffix
// is stripped: the literal suffix ends glob+suf, so glob+suf matches name
// exactly when name ends in suf and glob matches the rest.
func (g *globIndex) hits(name string) []globHit {
	var out []globHit
	try := func(n string) {
		out = append(out, g.exact[n]...)
		for i := range len(n) {
			if n[i] == '.' {
				out = append(out, g.ext[n[i:]]...)
			}
		}
		for _, r := range g.other {
			if ok, _ := filepath.Match(r.pattern, n); ok {
				out = append(out, r.hit)
			}
		}
	}
	try(name)
	for _, suf := range ignoredSuffixes {
		if rest, ok := strings.CutSuffix(name, suf); ok {
			try(rest)
		}
	}
	return out
}

type builtinIndex struct {
	lexers  []chroma.Lexer
	byName  map[string]chroma.Lexer
	byAlias map[string]chroma.Lexer
	tiers   [2]globIndex // Filenames, then AliasFilenames
}

var builtins = sync.OnceValue(func() *builtinIndex {
	ix := &builtinIndex{
		lexers:  lexers.GlobalLexerRegistry.Lexers,
		byName:  map[string]chroma.Lexer{},
		byAlias: map[string]chroma.Lexer{},
	}
	for t := range ix.tiers {
		ix.tiers[t] = globIndex{exact: map[string][]globHit{}, ext: map[string][]globHit{}}
	}
	for i, lx := range ix.lexers {
		cfg := lx.Config()
		ix.byName[cfg.Name] = lx
		ix.byName[strings.ToLower(cfg.Name)] = lx
		for _, a := range cfg.Aliases {
			ix.byAlias[a] = lx
			ix.byAlias[strings.ToLower(a)] = lx
		}
		for j, g := range cfg.Filenames {
			ix.tiers[0].add(g, globHit{i, j})
		}
		for j, g := range cfg.AliasFilenames {
			ix.tiers[1].add(g, globHit{i, j})
		}
	}
	return ix
})

// match is lexers.Match: the highest-priority lexer whose Filenames match,
// else whose AliasFilenames match, else nil.
func (ix *builtinIndex) match(filename string) chroma.Lexer {
	name := filepath.Base(filename)
	for t := range ix.tiers {
		hits := ix.tiers[t].hits(name)
		if len(hits) == 0 {
			continue
		}
		slices.SortFunc(hits, func(a, b globHit) int {
			return cmp.Or(cmp.Compare(a.lexer, b.lexer), cmp.Compare(a.glob, b.glob))
		})
		var matched chroma.PrioritisedLexers
		for _, h := range slices.Compact(hits) {
			matched = append(matched, ix.lexers[h.lexer])
		}
		sort.Sort(matched)
		return matched[0]
	}
	return nil
}

// get is lexers.Get: by name, by alias, then as a file extension or name.
func (ix *builtinIndex) get(name string) chroma.Lexer {
	lower := strings.ToLower(name)
	for _, lx := range []chroma.Lexer{ix.byName[name], ix.byAlias[name], ix.byName[lower], ix.byAlias[lower]} {
		if lx != nil {
			return lx
		}
	}
	var candidates chroma.PrioritisedLexers
	for _, f := range []string{"filename." + name, name} {
		if lx := ix.match(f); lx != nil {
			candidates = append(candidates, lx)
		}
	}
	if len(candidates) == 0 {
		return nil
	}
	sort.Sort(candidates)
	return candidates[0]
}

// names returns every lowercase name get resolves without wildcard
// globs: lexer names and aliases, extensions and literal filenames.
func (ix *builtinIndex) names() []string {
	var out []string
	for k := range ix.byName {
		out = append(out, strings.ToLower(k))
	}
	for k := range ix.byAlias {
		out = append(out, strings.ToLower(k))
	}
	for _, t := range ix.tiers {
		for k := range t.exact {
			out = append(out, strings.ToLower(k))
		}
		for k := range t.ext {
			out = append(out, strings.ToLower(k[1:]))
		}
	}
	return out
}