package tui

import (
	"strings"

	"github.com/alecthomas/chroma/v2"
	"github.com/alecthomas/chroma/v2/styles"
	"github.com/charmbracelet/glamour"
	glamourstyle "github.com/charmbracelet/glamour/ansi"
	"github.com/charmbracelet/lipgloss"
	"github.com/charmbracelet/x/ansi"
	"github.com/muesli/termenv"
)

// Assistant text is Markdown, so the transcript renders it as Markdown:
// headings become headings, ** becomes weight, a table becomes a table. The
// work is done by glamour, which is goldmark's CommonMark parser behind a
// terminal renderer; what this file owns is the stylesheet, which is derived
// from the theme in force so Markdown is in the same palette as everything
// else on the screen rather than in glamour's own.
//
// Only assistant text goes through here. Tool results, local command output
// and notices are program output, not prose, and are shown as they came.

// minMarkdownWidth is the point below which layout stops being worth
// attempting: a table or an indented code block in fifteen columns is noise,
// and plain wrapped text is the more useful thing to show.
const minMarkdownWidth = 20

// mdRenderers holds one renderer per width and colour depth: constructing one
// builds a goldmark pipeline and bakes the stylesheet in, and the transcript
// asks for a render on every streamed token.
var mdRenderers map[mdKey]*glamour.TermRenderer

// mdKey is everything a built renderer has baked in. The colour profile is
// part of it because a test that forces colour on, or a terminal that reports
// its depth late, must not be served a renderer built for the old answer.
type mdKey struct {
	width   int
	profile termenv.Profile
}

// renderMarkdown lays out assistant prose to fit width. It never fails in a
// way the reader can see: anything glamour cannot handle falls back to the
// plain wrap the transcript used before, because a message that renders badly
// is still better than a message that does not render.
func renderMarkdown(src string, width int) string {
	src = strings.TrimRight(src, "\n")
	if strings.TrimSpace(src) == "" {
		return ""
	}
	if width < minMarkdownWidth {
		return plainWrap(src, width)
	}
	r := markdownRenderer(width)
	if r == nil {
		return plainWrap(src, width)
	}
	out, err := r.Render(src)
	if err != nil {
		return plainWrap(src, width)
	}
	// glamour frames a document in blank lines. The transcript already
	// separates entries, so they would read as a gap that means nothing.
	return fit(strings.Trim(out, "\n"), width)
}

// fit squares glamour's output up with the columns we asked for.
//
// Two things need doing per line. The padding glamour lays down out to the
// wrap width comes off: it is invisible until someone selects the text, and
// then every line has a ragged tail of highlighted nothing. And a word too
// long to wrap — a URL, an import path, a hash — gets broken by force, because
// word wrapping leaves it hanging over the edge and the viewport then wraps
// it again at a column this layout knows nothing about.
func fit(s string, width int) string {
	var out []string
	for _, line := range strings.Split(s, "\n") {
		bare := strings.TrimRight(ansi.Strip(line), " \t")
		line = ansi.Truncate(line, ansi.StringWidth(bare), "")
		if ansi.StringWidth(line) <= width {
			out = append(out, compact(line))
			continue
		}
		for _, w := range strings.Split(ansi.Hardwrap(line, width, false), "\n") {
			out = append(out, compact(w))
		}
	}
	return strings.Join(out, "\n")
}

// compact drops the escape sequences in a line that paint nothing.
//
// glamour styles each word and each cell of padding separately, so one line of
// prose arrives as the same colour set and cleared thirty times over, followed
// by an escape per trimmed padding space. Stripping to the visible width
// leaves all of that behind — for a screenful of answer it is an order of
// magnitude more bytes of escape than of text, held for the whole session and
// re-scanned by the viewport on every frame.
//
// What survives is every escape that precedes text it actually changes, in
// order, and one reset to close the line.
func compact(line string) string {
	var (
		b       strings.Builder
		pending []string // escapes seen since the last visible text
		active  []string // what is in force on screen
		reset   bool     // a reset was seen and has not been emitted yet
	)
	for _, tok := range sgrTokens(line) {
		switch {
		case tok == "\x1b[0m" || tok == "\x1b[m":
			pending, reset = nil, len(active) > 0
		case strings.HasPrefix(tok, "\x1b[") && strings.HasSuffix(tok, "m"):
			pending = append(pending, tok)
		default:
			if reset {
				b.WriteString("\x1b[0m")
				active, reset = nil, false
			}
			for _, esc := range pending {
				b.WriteString(esc)
			}
			active = append(active, pending...)
			pending = nil
			b.WriteString(tok)
		}
	}
	if len(active) > 0 {
		b.WriteString("\x1b[0m")
	}
	return b.String()
}

// sgrTokens splits a line into SGR sequences and the text between them.
// Anything else escaped — a hyperlink, a cursor move — comes back as text, so
// it is passed through untouched rather than reasoned about.
func sgrTokens(s string) []string {
	var out []string
	text := 0
	for i := 0; i < len(s); {
		if s[i] != 0x1b || i+1 >= len(s) || s[i+1] != '[' {
			i++
			continue
		}
		j := i + 2
		for j < len(s) && (s[j] == ';' || (s[j] >= '0' && s[j] <= '9')) {
			j++
		}
		if j >= len(s) || s[j] != 'm' {
			i++
			continue
		}
		if text < i {
			out = append(out, s[text:i])
		}
		out = append(out, s[i:j+1])
		i, text = j+1, j+1
	}
	if text < len(s) {
		out = append(out, s[text:])
	}
	return out
}

// plainWrap is what the transcript did before Markdown, kept as the fallback.
func plainWrap(src string, width int) string {
	return styText.Width(maxInt(1, width)).Render(src)
}

func markdownRenderer(width int) *glamour.TermRenderer {
	// lipgloss's profile rather than termenv's own: glamour assumes truecolor,
	// and the rest of this UI paints with whatever lipgloss decided the
	// terminal can show. Two answers to that question on one screen is how
	// half a transcript ends up in colour and half of it not.
	key := mdKey{width: width, profile: lipgloss.ColorProfile()}
	if r, ok := mdRenderers[key]; ok {
		return r
	}
	r, err := glamour.NewTermRenderer(
		glamour.WithStyles(markdownStyle(currentTheme, key.profile, width)),
		glamour.WithWordWrap(width),
		glamour.WithColorProfile(key.profile),
		glamour.WithChromaFormatter(chromaFormatter(key.profile)),
	)
	if err != nil {
		return nil
	}
	if mdRenderers == nil {
		mdRenderers = map[mdKey]*glamour.TermRenderer{}
	}
	mdRenderers[key] = r
	return r
}

// resetMarkdown drops the renderers built against the previous palette, so the
// next draw builds them from the new one. applyTheme calls it.
func resetMarkdown() { mdRenderers = nil }

// chromaFormatter names the chroma writer that matches what the terminal can
// show. glamour defaults to terminal256, which would quantise this theme's
// colours to the nearest xterm slot and leave highlighted code a shade off
// every other thing on the screen.
func chromaFormatter(p termenv.Profile) string {
	switch p {
	case termenv.TrueColor:
		return "terminal16m"
	case termenv.ANSI:
		return "terminal"
	default:
		return "terminal256"
	}
}

// chromaThemes records which palettes have been handed to chroma's global
// registry, which is keyed by name and shared by every user of the package.
var chromaThemes = map[string]bool{}

// chromaTheme registers this theme's syntax colours with chroma and returns
// the name to ask for. It returns "" on a terminal that cannot show colour,
// which is how a code block is told to render plain: glamour only guards its
// own stylesheet against that case, not a theme named directly.
//
// The name is per-theme because chroma's registry is global and write-once in
// practice — glamour registers its own under the single name "charm" and skips
// the work if that name is taken, so a shared name would freeze code blocks in
// whichever palette happened to draw first.
func chromaTheme(t theme, p termenv.Profile) string {
	if p == termenv.Ascii {
		return ""
	}
	name := "harness-" + t.name
	if chromaThemes[name] {
		return name
	}
	c := markdownChroma(t)
	styles.Register(chroma.MustNewStyle(name, chroma.StyleEntries{
		chroma.Text:                chromaEntry(c.Text),
		chroma.Error:               chromaEntry(c.Error),
		chroma.Comment:             chromaEntry(c.Comment),
		chroma.CommentPreproc:      chromaEntry(c.CommentPreproc),
		chroma.Keyword:             chromaEntry(c.Keyword),
		chroma.KeywordReserved:     chromaEntry(c.KeywordReserved),
		chroma.KeywordNamespace:    chromaEntry(c.KeywordNamespace),
		chroma.KeywordType:         chromaEntry(c.KeywordType),
		chroma.Operator:            chromaEntry(c.Operator),
		chroma.Punctuation:         chromaEntry(c.Punctuation),
		chroma.Name:                chromaEntry(c.Name),
		chroma.NameBuiltin:         chromaEntry(c.NameBuiltin),
		chroma.NameTag:             chromaEntry(c.NameTag),
		chroma.NameAttribute:       chromaEntry(c.NameAttribute),
		chroma.NameClass:           chromaEntry(c.NameClass),
		chroma.NameConstant:        chromaEntry(c.NameConstant),
		chroma.NameDecorator:       chromaEntry(c.NameDecorator),
		chroma.NameException:       chromaEntry(c.NameException),
		chroma.NameFunction:        chromaEntry(c.NameFunction),
		chroma.NameOther:           chromaEntry(c.NameOther),
		chroma.Literal:             chromaEntry(c.Literal),
		chroma.LiteralNumber:       chromaEntry(c.LiteralNumber),
		chroma.LiteralDate:         chromaEntry(c.LiteralDate),
		chroma.LiteralString:       chromaEntry(c.LiteralString),
		chroma.LiteralStringEscape: chromaEntry(c.LiteralStringEscape),
		chroma.GenericDeleted:      chromaEntry(c.GenericDeleted),
		chroma.GenericInserted:     chromaEntry(c.GenericInserted),
		chroma.GenericEmph:         chromaEntry(c.GenericEmph),
		chroma.GenericStrong:       chromaEntry(c.GenericStrong),
		chroma.GenericSubheading:   chromaEntry(c.GenericSubheading),
	}))
	chromaThemes[name] = true
	return name
}

// markdownChroma maps chroma's token classes onto the same theme roles the
// rest of the UI uses, so highlighted code reads as part of this theme and not
// as a photograph of somebody else's editor.
func markdownChroma(t theme) glamourstyle.Chroma {
	var (
		text   = colorOf(t.text)
		dim    = colorOf(t.dim)
		accent = colorOf(t.accent)
		info   = colorOf(t.info)
		ok     = colorOf(t.ok)
		warn   = colorOf(t.warn)
		bad    = colorOf(t.err)
	)
	yes := ptr(true)
	return glamourstyle.Chroma{
		Text:                glamourstyle.StylePrimitive{Color: text},
		Error:               glamourstyle.StylePrimitive{Color: bad},
		Comment:             glamourstyle.StylePrimitive{Color: dim, Italic: yes},
		CommentPreproc:      glamourstyle.StylePrimitive{Color: warn},
		Keyword:             glamourstyle.StylePrimitive{Color: accent},
		KeywordReserved:     glamourstyle.StylePrimitive{Color: accent},
		KeywordNamespace:    glamourstyle.StylePrimitive{Color: accent},
		KeywordType:         glamourstyle.StylePrimitive{Color: info},
		Operator:            glamourstyle.StylePrimitive{Color: dim},
		Punctuation:         glamourstyle.StylePrimitive{Color: dim},
		Name:                glamourstyle.StylePrimitive{Color: text},
		NameBuiltin:         glamourstyle.StylePrimitive{Color: info},
		NameTag:             glamourstyle.StylePrimitive{Color: accent},
		NameAttribute:       glamourstyle.StylePrimitive{Color: info},
		NameClass:           glamourstyle.StylePrimitive{Color: info, Bold: yes},
		NameConstant:        glamourstyle.StylePrimitive{Color: warn},
		NameDecorator:       glamourstyle.StylePrimitive{Color: warn},
		NameException:       glamourstyle.StylePrimitive{Color: bad},
		NameFunction:        glamourstyle.StylePrimitive{Color: info},
		NameOther:           glamourstyle.StylePrimitive{Color: text},
		Literal:             glamourstyle.StylePrimitive{Color: ok},
		LiteralNumber:       glamourstyle.StylePrimitive{Color: warn},
		LiteralDate:         glamourstyle.StylePrimitive{Color: warn},
		LiteralString:       glamourstyle.StylePrimitive{Color: ok},
		LiteralStringEscape: glamourstyle.StylePrimitive{Color: warn},
		GenericDeleted:      glamourstyle.StylePrimitive{Color: bad},
		GenericInserted:     glamourstyle.StylePrimitive{Color: ok},
		GenericEmph:         glamourstyle.StylePrimitive{Italic: yes},
		GenericStrong:       glamourstyle.StylePrimitive{Bold: yes},
		GenericSubheading:   glamourstyle.StylePrimitive{Color: dim},
	}
}

// chromaEntry writes one token style in the "#rrggbb italic bold" grammar
// chroma parses its stylesheets from.
func chromaEntry(s glamourstyle.StylePrimitive) string {
	var parts []string
	if s.Color != nil {
		parts = append(parts, *s.Color)
	}
	if s.Italic != nil && *s.Italic {
		parts = append(parts, "italic")
	}
	if s.Bold != nil && *s.Bold {
		parts = append(parts, "bold")
	}
	return strings.Join(parts, " ")
}

// markdownStyle maps Markdown onto the ten theme roles.
//
// Headings carry no hashes. A hash is an instruction to make the line a
// heading, and having obeyed it there is no reason to print it too; the levels
// are told apart by colour and weight, the way the rest of this UI tells
// things apart.
func markdownStyle(t theme, profile termenv.Profile, width int) glamourstyle.StyleConfig {
	var (
		text   = colorOf(t.text)
		dim    = colorOf(t.dim)
		faint  = colorOf(t.faint)
		accent = colorOf(t.accent)
		info   = colorOf(t.info)
	)
	no, yes := ptr(false), ptr(true)

	return glamourstyle.StyleConfig{
		Document: glamourstyle.StyleBlock{
			StylePrimitive: glamourstyle.StylePrimitive{Color: text},
			// The transcript owns the left margin: everything hangs off the
			// dot in the gutter, and a second indent from glamour would put
			// prose out of line with every other kind of entry.
			Margin: ptr(uint(0)),
		},
		// Paragraph and Text are deliberately left bare. glamour applies
		// Styles.Text to every text run as an override, so a colour here would
		// win over the heading, the link and the blockquote alike and flatten
		// the whole document back to one colour. Document above is the base;
		// everything else is a departure from it.
		Paragraph: glamourstyle.StyleBlock{},

		Heading: glamourstyle.StyleBlock{StylePrimitive: glamourstyle.StylePrimitive{
			Color: accent, Bold: yes, BlockSuffix: "\n",
		}},
		H1: glamourstyle.StyleBlock{StylePrimitive: glamourstyle.StylePrimitive{Color: accent, Bold: yes}},
		H2: glamourstyle.StyleBlock{StylePrimitive: glamourstyle.StylePrimitive{Color: accent, Bold: yes}},
		H3: glamourstyle.StyleBlock{StylePrimitive: glamourstyle.StylePrimitive{Color: text, Bold: yes}},
		H4: glamourstyle.StyleBlock{StylePrimitive: glamourstyle.StylePrimitive{Color: text, Bold: yes}},
		H5: glamourstyle.StyleBlock{StylePrimitive: glamourstyle.StylePrimitive{Color: dim, Bold: yes}},
		H6: glamourstyle.StyleBlock{StylePrimitive: glamourstyle.StylePrimitive{Color: dim, Bold: no}},

		Strong:        glamourstyle.StylePrimitive{Bold: yes},
		Emph:          glamourstyle.StylePrimitive{Italic: yes},
		Strikethrough: glamourstyle.StylePrimitive{CrossedOut: yes, Color: dim},

		BlockQuote: glamourstyle.StyleBlock{
			StylePrimitive: glamourstyle.StylePrimitive{Color: dim, Italic: yes},
			Indent:         ptr(uint(1)),
			IndentToken:    ptr("│ "),
		},

		List:        glamourstyle.StyleList{LevelIndent: 2},
		Item:        glamourstyle.StylePrimitive{BlockPrefix: "• "},
		Enumeration: glamourstyle.StylePrimitive{BlockPrefix: ". "},
		Task:        glamourstyle.StyleTask{Ticked: "[✓] ", Unticked: "[ ] "},

		Link:     glamourstyle.StylePrimitive{Color: info, Underline: yes},
		LinkText: glamourstyle.StylePrimitive{Color: text},
		Image:    glamourstyle.StylePrimitive{Color: info, Underline: yes},
		ImageText: glamourstyle.StylePrimitive{
			Color: dim, Format: "Image: {{.text}} →",
		},

		// A rule divides the page, so it is drawn as wide as the page. The
		// stylesheet can say that because it is built per width, the same as
		// the renderer that holds it.
		HorizontalRule: glamourstyle.StylePrimitive{
			Color: faint, Format: "\n" + strings.Repeat("─", maxInt(1, width)) + "\n",
		},

		// Inline code takes the accent so a literal token stands out of the
		// sentence around it. A background colour would be the obvious other
		// choice, and is what glamour ships, but it tears across a wrap.
		Code: glamourstyle.StyleBlock{StylePrimitive: glamourstyle.StylePrimitive{Color: accent}},

		CodeBlock: glamourstyle.StyleCodeBlock{
			StyleBlock: glamourstyle.StyleBlock{
				StylePrimitive: glamourstyle.StylePrimitive{Color: text},
				Margin:         ptr(uint(2)),
			},
			// Named rather than handed over inline: glamour would register an
			// inline Chroma block under one global name shared with every
			// other program using it. See chromaTheme.
			Theme: chromaTheme(t, profile),
		},

		// The rules take the dim grey the rest of the UI draws structure in,
		// leaving the cells at the document's own colour.
		Table: glamourstyle.StyleTable{
			StyleBlock:      glamourstyle.StyleBlock{StylePrimitive: glamourstyle.StylePrimitive{Color: dim}},
			CenterSeparator: ptr("┼"),
			ColumnSeparator: ptr("│"),
			RowSeparator:    ptr("─"),
		},

		DefinitionTerm:        glamourstyle.StylePrimitive{Bold: yes},
		DefinitionDescription: glamourstyle.StylePrimitive{BlockPrefix: "\n  "},
	}
}

// colorOf resolves an adaptive colour to the one hex value glamour can take.
// glamour's stylesheet is strings, so the light/dark choice has to be made
// here rather than at draw time; it is remade whenever the theme is applied,
// which is the only moment the answer can change.
func colorOf(c lipgloss.AdaptiveColor) *string {
	s := c.Light
	if lipgloss.HasDarkBackground() {
		s = c.Dark
	}
	return &s
}

func ptr[T any](v T) *T { return &v }