package tty

import (
	"io"
	"strings"

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

type ChromaHighlighter struct {
	lexer     chroma.Lexer
	formatter chroma.Formatter
	style     *chroma.Style
}

func NewChromaHighlighter(lang, format, style string) *ChromaHighlighter {
	l := lexers.Get(lang)
	if l == nil {
		l = lexers.Fallback
	}
	l = chroma.Coalesce(l)

	f := formatters.Get(format)
	if f == nil {
		f = formatters.Fallback
	}

	s := styles.Get(style)
	if s == nil {
		s = styles.Fallback
	}

	return &ChromaHighlighter{
		lexer:     l,
		formatter: f,
		style:     s,
	}
}

func (s *ChromaHighlighter) Highlight(w io.Writer, text string) error {
	it, err := s.lexer.Tokenise(nil, text)
	if err != nil {
		return err
	}
	return s.formatter.Format(w, s.style, it)
}

func (s *ChromaHighlighter) HighlightS(text string) (string, error) {
	it, err := s.lexer.Tokenise(nil, text)
	if err != nil {
		return "", err
	}
	sb := strings.Builder{}
	sb.Grow(len(text) * 2)
	s.formatter.Format(&sb, s.style, it)
	return sb.String(), nil
}