package tui

import (
	"strings"

	"github.com/charmbracelet/x/ansi"
)

// markdown renders the Markdown models write, enough of it to read well in
// a terminal: headings, lists, quotes, rules, fenced code and tables, and
// bold, italic, strikethrough, code and links within a line. Lines stay as
// they are rather than being joined into paragraphs, since models break
// lines where they mean to. s must already be sanitized: the only escape
// sequences in what this returns are the styles it adds.
func markdown(s string, width int) string {
	width = max(width, 10)
	lines := strings.Split(s, "\n")
	var out []string
	for i := 0; i < len(lines); i++ {
		line := lines[i]
		trimmed := strings.TrimSpace(line)
		switch {
		case fenceOf(trimmed) != "":
			// Until the closing fence, or the end of a reply still
			// streaming in. The info string (a language, usually) shows
			// too: nothing the model wrote is hidden.
			fence := fenceOf(trimmed)
			if info := strings.TrimLeft(trimmed, fence[:1]); strings.TrimSpace(info) != "" {
				out = append(out, dimStyle.Render("  "+strings.TrimSpace(info)))
			}
			for i++; i < len(lines) && !closesFence(strings.TrimSpace(lines[i]), fence); i++ {
				for _, c := range strings.Split(ansi.Hardwrap(lines[i], width-2, true), "\n") {
					out = append(out, "  "+style{code: true}.sgr()+c+sgrReset)
				}
			}
		case tableStart(lines, i):
			aligns := alignments(lines[i+1])
			rows := [][]string{cells(line)}
			for i += 2; i < len(lines) && strings.Contains(lines[i], "|"); i++ {
				rows = append(rows, cells(lines[i]))
			}
			i--
			out = append(out, table(rows, aligns, width))
		case isRule(trimmed):
			out = append(out, dimStyle.Render(strings.Repeat("─", width)))
		case headingLevel(trimmed) > 0:
			n := headingLevel(trimmed)
			base := "1;94" // bold, in the accent color (ANSI 12)
			if n == 1 {
				base += ";4"
			}
			out = append(out, carry(ansi.Wrap(inline(strings.TrimSpace(trimmed[n:]), base), width, "")))
		case strings.HasPrefix(trimmed, ">"):
			quote := strings.TrimPrefix(strings.TrimPrefix(trimmed, ">"), " ")
			out = append(out, indent(carry(ansi.Wrap(inline(quote, "3"), width-2, "")), dimStyle.Render("│ ")))
		default:
			if lead, marker, text, ok := listItem(line); ok {
				prefix := strings.Repeat(" ", lead) + marker + " "
				w := ansi.StringWidth(prefix)
				body := carry(ansi.Wrap(inline(text, ""), width-w, ""))
				out = append(out, prefix+strings.ReplaceAll(body, "\n", "\n"+strings.Repeat(" ", w)))
				continue
			}
			out = append(out, carry(ansi.Wrap(inline(line, ""), width, "")))
		}
	}
	return strings.Join(out, "\n")
}

const sgrReset = "\x1b[0m"

// style is the styling at a point in a line.
type style struct {
	base                             string // SGR parameters for the whole line, such as "1" for bold
	bold, italic, strike, code, link bool
	mark                             bool // Markdown syntax, dimmed in the input
	sel                              bool // selected in the input
}

// markSGR colors Markdown syntax: like dimStyle, darker on a light
// background (see the BackgroundColorMsg case in Update).
var markSGR = "38;5;250"

// sgr is one escape sequence that sets exactly this style, so the last
// sequence before any point in a line says what style it has there.
func (st style) sgr() string {
	p := []string{"0"}
	if st.base != "" {
		p = append(p, st.base)
	}
	for _, a := range []struct {
		on   bool
		code string
	}{{st.bold, "1"}, {st.italic, "3"}, {st.strike, "9"}, {st.code, "36"}, {st.link, "4"}, {st.mark, markSGR}, {st.sel, "7"}} {
		if a.on {
			p = append(p, a.code)
		}
	}
	return "\x1b[" + strings.Join(p, ";") + "m"
}

// carry makes each line of styled text stand alone: it starts with the
// style the line before ended in and ends with a reset, so a line can be
// shown without the ones before it, and styles do not run into borders.
func carry(s string) string {
	if !strings.Contains(s, "\x1b") {
		return s
	}
	lines := strings.Split(s, "\n")
	state := ""
	for i, l := range lines {
		l = state + l
		if j := strings.LastIndex(l, "\x1b["); j >= 0 {
			if k := strings.IndexByte(l[j:], 'm'); k >= 0 {
				state = l[j : j+k+1]
			}
		}
		lines[i] = l + sgrReset
	}
	return strings.Join(lines, "\n")
}

// span is a run of a line's raw text and how it shows.
type span struct {
	text string
	st   style
	kind byte // 0 for text, 'm' for Markdown syntax, 'u' for a link's URL
	// The construct, such as a code span, that the syntax is part of,
	// numbered from 1: the input shows a construct's syntax only while
	// the cursor is within it.
	group int
	// The syntax opens a construct that is not closed yet, which runs to
	// the end of the line.
	unclosed bool
}

// spans parses the Markdown within a line, with base the SGR parameters
// for all of it. Their text joined is the line. In the input, cursor is
// where in the line the cursor is (-1 elsewhere): syntax before it that
// opens a construct opens it even if not closed yet, so the text being
// typed shows styled.
func spans(s, base string, cursor int) []span {
	st := style{base: base}
	var out []span
	group, groups := 0, 0
	var open [3]int // the groups of the bold, italic and strikethrough open
	add := func(text string, kind byte) {
		switch {
		case text == "":
		case kind == 0:
			out = append(out, span{text, st, kind, 0, false})
		case kind == 'm':
			out = append(out, span{text, style{base: base}, kind, group, false})
		default:
			out = append(out, span{text, st, kind, group, false})
		}
	}
	newGroup := func() {
		groups++
		group = groups
	}
	start := 0 // the start of plain text not yet added
	for i := 0; i < len(s); {
		c := s[i]
		switch {
		case c == '\\' && i+1 < len(s) && strings.IndexByte("\\`*_~[]()#|>-+.!", s[i+1]) >= 0:
			add(s[start:i], 0)
			newGroup()
			add(s[i:i+1], 'm')
			start = i + 1 // the escaped character starts the next text
			i += 2
			continue
		case c == '`':
			n := runOf(s[i:], '`')
			tick := s[i : i+n]
			end := strings.Index(s[i+n:], tick)
			if end < 0 && cursor >= i+n {
				add(s[start:i], 0)
				newGroup()
				add(tick, 'm')
				out[len(out)-1].unclosed = true
				st.code = true
				add(s[i+n:], 0)
				return out
			}
			if end < 0 {
				i += n
				continue
			}
			add(s[start:i], 0)
			code, opener, closer := s[i+n:i+n+end], tick, tick
			if len(code) > 2 && code[0] == ' ' && code[len(code)-1] == ' ' {
				code, opener, closer = code[1:len(code)-1], tick+" ", " "+tick
			}
			newGroup()
			add(opener, 'm')
			st.code = true
			add(code, 0)
			st.code = false
			add(closer, 'm')
			i += 2*n + end
			start = i
			continue
		case c == '[':
			if text, url, n, ok := link(s[i:]); ok {
				add(s[start:i], 0)
				newGroup()
				add("[", 'm')
				st.link = true
				add(text, 0)
				st.link = false
				if url == text {
					add("]("+url+")", 'm')
				} else {
					add("](", 'm')
					add(url, 'u')
					add(")", 'm')
				}
				i += n
				start = i
				continue
			}
		case c == '*' || c == '_' || c == '~':
			n := runOf(s[i:], c)
			if next := st; toggle(&next, s, i, n, cursor >= i+n) {
				add(s[start:i], 0)
				// The delimiters that close a style are in the group of
				// those that opened it, even when they differ, as *** and **.
				was := [3]bool{st.bold, st.italic, st.strike}
				now := [3]bool{next.bold, next.italic, next.strike}
				opening := now[0] && !was[0] || now[1] && !was[1] || now[2] && !was[2]
				if opening {
					newGroup()
				}
				for f := range open {
					switch {
					case now[f] && !was[f]:
						open[f] = group
					case was[f] && !now[f]:
						group = open[f]
					}
				}
				st = next
				add(s[i:i+n], 'm')
				out[len(out)-1].unclosed = opening && !strings.Contains(s[i+n:], s[i:i+n])
				start = i + n
			}
			i += n
			continue
		}
		i++
	}
	add(s[start:], 0)
	return out
}

// inline renders the Markdown within a line, with base the SGR parameters
// for all of it.
func inline(s, base string) string {
	var b strings.Builder
	cur := style{base: base}
	if base != "" {
		b.WriteString(cur.sgr())
	}
	for _, sp := range spans(s, base, -1) {
		text := sp.text
		switch sp.kind {
		case 'm':
			continue
		case 'u':
			text = " (" + text + ")"
		}
		if sp.st != cur {
			b.WriteString(sp.st.sgr())
			cur = sp.st
		}
		b.WriteString(text)
	}
	if strings.Contains(b.String(), "\x1b") {
		b.WriteString(sgrReset)
	}
	return b.String()
}

// toggle turns the style a run of n delimiters (*, _ or ~) at s[i] marks on
// or off, if it opens or closes one, and reports whether it did. With
// unclosed, it may open a style that nothing after it closes.
func toggle(st *style, s string, i, n int, unclosed bool) bool {
	c := s[i]
	var flags []*bool
	switch {
	case c == '~' && n == 2:
		flags = []*bool{&st.strike}
	case c == '~':
		return false
	case n == 1:
		flags = []*bool{&st.italic}
	case n == 2:
		flags = []*bool{&st.bold}
	case n == 3 && st.bold == st.italic:
		flags = []*bool{&st.bold, &st.italic}
	default:
		return false
	}
	var prev, next byte = ' ', ' '
	if i > 0 {
		prev = s[i-1]
	}
	if i+n < len(s) {
		next = s[i+n]
	}
	on := *flags[0]
	if on {
		// Closing: after text, and for _ not within a word.
		if prev == ' ' || c == '_' && isAlnum(next) {
			return false
		}
	} else {
		// Opening: before text that it closes after, and for _ not
		// within a word, so snake_case stays as it is.
		if next == ' ' || !unclosed && !strings.Contains(s[i+n:], s[i:i+n]) || c == '_' && isAlnum(prev) {
			return false
		}
	}
	for _, f := range flags {
		*f = !on
	}
	return true
}

// link parses [text](url) at the start of s, returning how long it is.
func link(s string) (text, url string, n int, ok bool) {
	mid := strings.Index(s, "](")
	if mid < 1 {
		return "", "", 0, false
	}
	end := strings.IndexByte(s[mid+2:], ')')
	if end < 1 {
		return "", "", 0, false
	}
	text, url = s[1:mid], s[mid+2:mid+2+end]
	if strings.ContainsAny(text, "[]") || strings.ContainsAny(url, " ()") {
		return "", "", 0, false
	}
	return text, url, mid + 3 + end, true
}

func runOf(s string, c byte) int {
	n := 0
	for n < len(s) && s[n] == c {
		n++
	}
	return n
}

func isAlnum(c byte) bool {
	return c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c >= '0' && c <= '9' || c >= 0x80
}

// headingLevel is the number of #s starting a heading, or 0 for a line
// that is not one.
func headingLevel(t string) int {
	n := runOf(t, '#')
	if n == 0 || n > 6 || n < len(t) && t[n] != ' ' {
		return 0
	}
	return n
}

// fenceOf returns the fence a trimmed line opens a code block with, ``` or
// ~~~, or "" if it opens none.
func fenceOf(t string) string {
	if strings.HasPrefix(t, "```") || strings.HasPrefix(t, "~~~") {
		return t[:3]
	}
	return ""
}

// closesFence reports whether a trimmed line closes a code block opened
// with fence: it is only fence characters. With more on the line, such as
// another block's info string, it is code.
func closesFence(t, fence string) bool {
	return strings.HasPrefix(t, fence) && strings.Trim(t, fence[:1]) == ""
}

// tableStart reports whether lines[i] is a table's header: a row with a
// delimiter row under it.
func tableStart(lines []string, i int) bool {
	return strings.Contains(lines[i], "|") && i+1 < len(lines) && isDelimiterRow(lines[i+1], len(cells(lines[i])))
}

// isRule reports whether a line is a thematic break: three or more of the
// same one of -, * or _, and maybe spaces.
func isRule(t string) bool {
	if len(t) < 3 || !strings.ContainsRune("-*_", rune(t[0])) {
		return false
	}
	n := 0
	for i := range len(t) {
		switch t[i] {
		case t[0]:
			n++
		case ' ':
		default:
			return false
		}
	}
	return n >= 3
}

// listItem parses a list item: its indentation, its marker (a bullet for
// -, * or +, or the number of a numbered one), and its text.
func listItem(l string) (lead int, marker, text string, ok bool) {
	n := listMarker(l)
	if n == 0 {
		return 0, "", "", false
	}
	lead = len(l) - len(strings.TrimLeft(l, " "))
	marker, text = l[lead:n-1], l[n:]
	if strings.IndexByte("-*+", marker[0]) < 0 {
		return lead, marker, text, true
	}
	switch {
	case strings.HasPrefix(text, "[ ] "):
		text = "☐ " + text[4:]
	case strings.HasPrefix(text, "[x] ") || strings.HasPrefix(text, "[X] "):
		text = "☑ " + text[4:]
	}
	return lead, "•", text, true
}

// listMarker is how long the start of a list item is: its indentation, its
// marker (-, * or +, or a number and . or ) after it) and a space. It is 0
// for a line that is not one.
func listMarker(l string) int {
	t := strings.TrimLeft(l, " ")
	lead := len(l) - len(t)
	if len(t) >= 2 && strings.IndexByte("-*+", t[0]) >= 0 && t[1] == ' ' {
		return lead + 2
	}
	n := 0
	for n < len(t) && n < 9 && t[n] >= '0' && t[n] <= '9' {
		n++
	}
	if n > 0 && n+1 < len(t) && (t[n] == '.' || t[n] == ')') && t[n+1] == ' ' {
		return lead + n + 2
	}
	return 0
}

// cells splits a table row at the |s not escaped or in code.
func cells(l string) []string {
	t := strings.TrimSpace(l)
	t = strings.TrimPrefix(t, "|")
	if strings.HasSuffix(t, "|") && !strings.HasSuffix(t, "\\|") {
		t = t[:len(t)-1]
	}
	var out []string
	start, code := 0, false
	for i := 0; i < len(t); i++ {
		switch t[i] {
		case '\\':
			i++
		case '`':
			code = !code
		case '|':
			if !code {
				out = append(out, strings.TrimSpace(t[start:i]))
				start = i + 1
			}
		}
	}
	return append(out, strings.TrimSpace(t[start:]))
}

// isDelimiterRow reports whether a line is the row under a table's header,
// such as |---|:--:|, for a header of n cells.
func isDelimiterRow(l string, n int) bool {
	if !strings.Contains(l, "-") {
		return false
	}
	cs := cells(l)
	if len(cs) != n {
		return false
	}
	for _, c := range cs {
		c = strings.TrimSuffix(strings.TrimPrefix(c, ":"), ":")
		if c == "" || strings.Trim(c, "-") != "" {
			return false
		}
	}
	return true
}

// alignments is how each column aligns, from the delimiter row: 'l', 'c'
// or 'r'.
func alignments(l string) []byte {
	var out []byte
	for _, c := range cells(l) {
		switch {
		case strings.HasPrefix(c, ":") && strings.HasSuffix(c, ":"):
			out = append(out, 'c')
		case strings.HasSuffix(c, ":"):
			out = append(out, 'r')
		default:
			out = append(out, 'l')
		}
	}
	return out
}

// table draws a table with borders, its header in bold, wrapping the text
// in its cells when it is too wide for the screen. A row with more cells
// than the header widens the table, rather than losing them as in GFM:
// nothing the model wrote may go unseen.
func table(rows [][]string, aligns []byte, width int) string {
	for _, row := range rows {
		for len(aligns) < len(row) {
			aligns = append(aligns, 'l')
		}
	}
	ncols := len(aligns)
	rendered := make([][]string, len(rows))
	widths := make([]int, ncols)
	for r, row := range rows {
		rendered[r] = make([]string, ncols)
		for j := range ncols {
			c := ""
			if j < len(row) {
				c = row[j]
			}
			base := ""
			if r == 0 {
				base = "1"
			}
			rendered[r][j] = inline(c, base)
			widths[j] = max(widths[j], ansi.StringWidth(rendered[r][j]), 1)
		}
	}
	// Each column takes its width plus " │ " between them, and "│ " and
	// " │" at the edges. Narrow the widest column until the table fits,
	// down to a few characters each.
	for {
		total, widest := 3*ncols+1, 0
		for j, w := range widths {
			total += w
			if w > widths[widest] {
				widest = j
			}
		}
		if total <= width || widths[widest] <= 3 {
			break
		}
		widths[widest]--
	}
	border := func(left, mid, right string) string {
		parts := make([]string, ncols)
		for j, w := range widths {
			parts[j] = strings.Repeat("─", w+2)
		}
		return dimStyle.Render(left + strings.Join(parts, mid) + right)
	}
	bar := dimStyle.Render("│")
	out := []string{border("┌", "┬", "┐")}
	for r, row := range rendered {
		wrapped := make([][]string, ncols)
		height := 1
		for j, c := range row {
			wrapped[j] = strings.Split(carry(ansi.Wrap(c, widths[j], "")), "\n")
			height = max(height, len(wrapped[j]))
		}
		for k := range height {
			var b strings.Builder
			b.WriteString(bar)
			for j := range ncols {
				c := ""
				if k < len(wrapped[j]) {
					c = wrapped[j][k]
				}
				pad := max(widths[j]-ansi.StringWidth(c), 0)
				left := 0
				switch aligns[j] {
				case 'r':
					left = pad
				case 'c':
					left = pad / 2
				}
				b.WriteString(" " + strings.Repeat(" ", left) + c + strings.Repeat(" ", pad-left) + " " + bar)
			}
			out = append(out, b.String())
		}
		if r == 0 {
			out = append(out, border("├", "┼", "┤"))
		}
	}
	out = append(out, border("└", "┴", "┘"))
	return strings.Join(out, "\n")
}