package chat

import (
	"strings"
	"unicode"
	"unicode/utf8"

	"github.com/alecthomas/chroma/v2/quick"
	"github.com/charmbracelet/lipgloss"
	"github.com/charmbracelet/x/ansi"
	"github.com/mattn/go-runewidth"
	"github.com/muesli/termenv"
)

func renderMarkdownForView(markdown string, width int) string {
	return renderMarkdownForViewWithCodeCache(markdown, width, nil)
}

type markdownCodeBlockCacheKey struct {
	language  string
	source    string
	formatter string
	style     string
}

func renderMarkdownForViewWithCodeCache(markdown string, width int, codeCache *map[markdownCodeBlockCacheKey]string) string {
	if width < 20 {
		width = 20
	}

	source := strings.Split(strings.TrimRight(markdown, "\n"), "\n")
	var rendered []string
	inCodeBlock := false
	codeLanguage := ""
	var codeLines []string
	for i := 0; i < len(source); i++ {
		line := strings.TrimRight(source[i], "\r")
		trimmed := strings.TrimSpace(line)

		if strings.HasPrefix(trimmed, "```") {
			if inCodeBlock {
				rendered = append(rendered, renderMarkdownCodeBlock(codeLines, codeLanguage, width, codeCache, true)...)
				codeLines = nil
				inCodeBlock = false
				codeLanguage = ""
			} else {
				inCodeBlock = true
				codeLanguage = markdownCodeFenceLanguage(trimmed)
			}
			continue
		}
		if inCodeBlock {
			codeLines = append(codeLines, line)
			continue
		}

		if table, consumed := renderMarkdownTable(source[i:], width); consumed > 0 {
			rendered = append(rendered, table...)
			i += consumed - 1
			continue
		}

		if heading, ok := markdownHeading(trimmed); ok {
			rendered = append(rendered, chatHeaderStyle.Render(renderMarkdownRunes(parseMarkdownInline(heading))))
			continue
		}

		if trimmed == "" {
			rendered = append(rendered, "")
			continue
		}
		rendered = append(rendered, wrapMarkdownInline(line, width)...)
	}
	if inCodeBlock {
		rendered = append(rendered, renderMarkdownCodeBlock(codeLines, codeLanguage, width, codeCache, false)...)
	}
	return strings.Join(rendered, "\n")
}

func splitRenderedBody(body string) []string {
	body = strings.TrimRight(body, "\n")
	if body == "" {
		return []string{""}
	}
	return strings.Split(body, "\n")
}

func markdownHeading(line string) (string, bool) {
	if !strings.HasPrefix(line, "#") {
		return "", false
	}
	level := 0
	for level < len(line) && line[level] == '#' {
		level++
	}
	if level == 0 || level > 6 || level >= len(line) || line[level] != ' ' {
		return "", false
	}
	return strings.TrimSpace(line[level:]), true
}

type markdownInlineStyle uint8

const (
	markdownPlain markdownInlineStyle = iota
	markdownStrong
	markdownCode
)

type markdownInlineRune struct {
	r     rune
	style markdownInlineStyle
}

// wrapMarkdownInline parses a complete source line before wrapping it. That
// keeps emphasis intact when its opening and closing delimiters land on
// different visual lines.
func wrapMarkdownInline(line string, width int) []string {
	return wrapInlineRunes(parseMarkdownInline(line), width)
}

func wrapInlineRunes(runes []markdownInlineRune, width int) []string {
	if len(runes) == 0 {
		return []string{""}
	}

	var rendered []string
	for len(runes) > 0 {
		hardCut, spaceCut, currentWidth := 0, 0, 0
		for i, item := range runes {
			nextWidth := currentWidth + runewidth.RuneWidth(item.r)
			if nextWidth > width {
				break
			}
			currentWidth = nextWidth
			hardCut = i + 1
			if unicode.IsSpace(item.r) && currentWidth > width/2 {
				spaceCut = i
			}
		}
		cut := hardCut
		if spaceCut > 0 {
			cut = spaceCut
		}
		if cut == 0 {
			cut = 1
		}

		lineRunes := trimMarkdownSpace(runes[:cut])
		rendered = append(rendered, renderMarkdownRunes(lineRunes))
		runes = trimMarkdownSpace(runes[cut:])
	}
	return rendered
}

func parseMarkdownInline(line string) []markdownInlineRune {
	var out []markdownInlineRune
	for len(line) > 0 {
		if strings.HasPrefix(line, "`") {
			if end := strings.Index(line[1:], "`"); end >= 0 {
				out = appendMarkdownRunes(out, line[1:end+1], markdownCode)
				line = line[end+2:]
				continue
			}
		}
		if (strings.HasPrefix(line, "**") || strings.HasPrefix(line, "__")) && canOpenMarkdownStrong(out) {
			delimiter := line[:2]
			if end := strings.Index(line[2:], delimiter); end >= 0 {
				out = appendMarkdownRunes(out, line[2:end+2], markdownStrong)
				line = line[end+4:]
				continue
			}
		}

		r, size := utf8.DecodeRuneInString(line)
		out = append(out, markdownInlineRune{r: r, style: markdownPlain})
		line = line[size:]
	}
	return out
}

// canOpenMarkdownStrong keeps delimiter-like text in bare URLs and identifiers
// literal, only treating ** / __ as strong emphasis at the common
// whitespace- or punctuation-delimited form.
func canOpenMarkdownStrong(out []markdownInlineRune) bool {
	if len(out) == 0 {
		return true
	}
	previous := out[len(out)-1].r
	return (unicode.IsSpace(previous) || unicode.IsPunct(previous)) && !markdownStrongInURL(out)
}

func markdownStrongInURL(out []markdownInlineRune) bool {
	start := len(out)
	for start > 0 && !unicode.IsSpace(out[start-1].r) {
		start--
	}

	var token strings.Builder
	for _, item := range out[start:] {
		token.WriteRune(item.r)
	}
	return strings.Contains(token.String(), "://")
}

func appendMarkdownRunes(out []markdownInlineRune, text string, style markdownInlineStyle) []markdownInlineRune {
	for _, r := range text {
		out = append(out, markdownInlineRune{r: r, style: style})
	}
	return out
}

func trimMarkdownSpace(runes []markdownInlineRune) []markdownInlineRune {
	start, end := 0, len(runes)
	for start < end && unicode.IsSpace(runes[start].r) {
		start++
	}
	for end > start && unicode.IsSpace(runes[end-1].r) {
		end--
	}
	return runes[start:end]
}

func renderMarkdownRunes(runes []markdownInlineRune) string {
	var b strings.Builder
	for start := 0; start < len(runes); {
		end := start + 1
		for end < len(runes) && runes[end].style == runes[start].style {
			end++
		}
		var text strings.Builder
		for _, item := range runes[start:end] {
			text.WriteRune(item.r)
		}
		switch runes[start].style {
		case markdownStrong:
			b.WriteString(chatStrongStyle.Render(text.String()))
		case markdownCode:
			b.WriteString(chatInlineCodeStyle.Render(text.String()))
		default:
			b.WriteString(text.String())
		}
		start = end
	}
	return b.String()
}

func renderMarkdownCodeBlock(source []string, language string, width int, codeCache *map[markdownCodeBlockCacheKey]string, complete bool) []string {
	codeWidth := max(1, width-2)
	code := strings.Join(source, "\n")
	lines := wrapChatText(code, codeWidth)
	if highlighted, ok := highlightMarkdownCodeBlock(language, code, codeCache, complete); ok {
		lines = wrapHighlightedMarkdownCode(highlighted, codeWidth)
	}
	for i, wrapped := range lines {
		lines[i] = "  " + chatCodeBlockStyle.Render(wrapped)
	}
	return lines
}

func markdownCodeFenceLanguage(fence string) string {
	info := strings.TrimSpace(strings.TrimPrefix(fence, "```"))
	if info == "" {
		return ""
	}
	return strings.Fields(info)[0]
}

func highlightMarkdownCodeBlock(language, source string, codeCache *map[markdownCodeBlockCacheKey]string, complete bool) (string, bool) {
	if language == "" || source == "" || lipgloss.ColorProfile() == termenv.Ascii {
		return source, false
	}

	key := markdownCodeBlockCacheKey{
		language:  language,
		source:    source,
		formatter: markdownCodeFormatter(),
		style:     markdownCodeStyle(),
	}
	if complete && codeCache != nil && *codeCache != nil {
		if highlighted, ok := (*codeCache)[key]; ok {
			return highlighted, true
		}
	}

	var rendered strings.Builder
	if err := quick.Highlight(&rendered, source, language, key.formatter, key.style); err != nil {
		return source, false
	}
	highlighted := strings.TrimSuffix(rendered.String(), "\n")
	if complete && codeCache != nil {
		if *codeCache == nil {
			*codeCache = make(map[markdownCodeBlockCacheKey]string)
		}
		(*codeCache)[key] = highlighted
	}
	return highlighted, true
}

func wrapHighlightedMarkdownCode(highlighted string, width int) []string {
	lines := strings.Split(ansi.Hardwrap(highlighted, width, true), "\n")
	activeStyle := ""
	for i, line := range lines {
		lines[i] = activeStyle + line
		activeStyle = markdownCodeANSIStyle(activeStyle, line)
		if activeStyle != "" && i < len(lines)-1 {
			lines[i] += "\x1b[0m"
		}
	}
	return lines
}

func markdownCodeANSIStyle(activeStyle, line string) string {
	for {
		start := strings.Index(line, "\x1b[")
		if start < 0 {
			return activeStyle
		}
		line = line[start+2:]
		end := strings.IndexByte(line, 'm')
		if end < 0 {
			return activeStyle
		}
		sequence := line[:end]
		if markdownCodeANSIReset(sequence) {
			activeStyle = ""
		} else {
			activeStyle += "\x1b[" + sequence + "m"
		}
		line = line[end+1:]
	}
}

func markdownCodeANSIReset(sequence string) bool {
	if sequence == "" {
		return true
	}
	for _, parameter := range strings.Split(sequence, ";") {
		if parameter == "0" {
			return true
		}
	}
	return false
}

func markdownCodeFormatter() string {
	switch lipgloss.ColorProfile() {
	case termenv.TrueColor:
		return "terminal16m"
	case termenv.ANSI256:
		return "terminal256"
	default:
		return "terminal16"
	}
}

func markdownCodeStyle() string {
	if lipgloss.HasDarkBackground() {
		return "github-dark"
	}
	return "github"
}

func renderMarkdownTable(lines []string, width int) ([]string, int) {
	if len(lines) < 2 || !looksLikeMarkdownTableRow(lines[0]) || !isMarkdownTableSeparator(lines[1]) {
		return nil, 0
	}

	var rows [][]string
	consumed := 0
	for consumed < len(lines) && looksLikeMarkdownTableRow(lines[consumed]) {
		if consumed == 1 && isMarkdownTableSeparator(lines[consumed]) {
			consumed++
			continue
		}
		rows = append(rows, parseMarkdownTableRow(lines[consumed]))
		consumed++
	}
	if len(rows) == 0 {
		return nil, 0
	}

	columnCount := 0
	for _, row := range rows {
		columnCount = max(columnCount, len(row))
	}
	naturalWidths := make([]int, columnCount)
	for _, row := range rows {
		for i := range columnCount {
			cell := ""
			if i < len(row) {
				cell = row[i]
			}
			naturalWidths[i] = max(naturalWidths[i], markdownInlineWidth(cell))
		}
	}
	widths := markdownTableColumnWidths(naturalWidths, width)

	var rendered []string
	for rowIndex, row := range rows {
		wrappedCells := make([][]string, columnCount)
		rowHeight := 1
		for i := range columnCount {
			cell := ""
			if i < len(row) {
				cell = row[i]
			}
			wrappedCells[i] = wrapMarkdownTableCell(cell, widths[i])
			rowHeight = max(rowHeight, len(wrappedCells[i]))
		}
		for lineIndex := range rowHeight {
			cells := make([]string, columnCount)
			for i := range columnCount {
				cellLine := ""
				if lineIndex < len(wrappedCells[i]) {
					cellLine = wrappedCells[i][lineIndex]
				}
				cells[i] = padPlainLine(cellLine, widths[i])
			}
			line := strings.Join(cells, chatTableBorderStyle.Render(" | "))
			if rowIndex == 0 {
				line = chatHeaderStyle.Render(stripANSIForWidth(line))
			}
			rendered = append(rendered, line)
		}
	}
	return rendered, consumed
}

func markdownTableColumnWidths(naturalWidths []int, width int) []int {
	if len(naturalWidths) == 0 {
		return nil
	}
	separatorWidth := max(0, len(naturalWidths)-1) * lipglossWidth(" | ")
	available := max(1, width-separatorWidth)
	widths := make([]int, len(naturalWidths))
	minWidths := make([]int, len(naturalWidths))
	for i, natural := range naturalWidths {
		widths[i] = max(1, natural)
		minWidth := min(widths[i], 12)
		if i == 0 {
			minWidth = min(widths[i], 4)
		}
		minWidths[i] = max(1, minWidth)
	}

	for sumInts(widths) > available {
		index := widestShrinkableColumn(widths, minWidths)
		if index < 0 {
			break
		}
		widths[index]--
	}
	for sumInts(widths) > available {
		index := widestColumn(widths)
		if index < 0 || widths[index] <= 1 {
			break
		}
		widths[index]--
	}
	return widths
}

func widestShrinkableColumn(widths, minWidths []int) int {
	index := -1
	for i, width := range widths {
		if width <= minWidths[i] {
			continue
		}
		if index < 0 || width > widths[index] {
			index = i
		}
	}
	return index
}

func widestColumn(widths []int) int {
	index := -1
	for i, width := range widths {
		if index < 0 || width > widths[index] {
			index = i
		}
	}
	return index
}

func sumInts(values []int) int {
	sum := 0
	for _, value := range values {
		sum += value
	}
	return sum
}

func wrapMarkdownTableCell(cell string, width int) []string {
	lines := wrapInlineRunes(parseMarkdownInline(cell), max(1, width))
	if len(lines) == 0 {
		return []string{""}
	}
	return lines
}

// markdownInlineWidth reports the visible width of a cell once Markdown
// delimiters are parsed away, so columns size to rendered content.
func markdownInlineWidth(cell string) int {
	width := 0
	for _, item := range parseMarkdownInline(cell) {
		width += runewidth.RuneWidth(item.r)
	}
	return width
}

func looksLikeMarkdownTableRow(line string) bool {
	line = strings.TrimSpace(line)
	return strings.Contains(line, "|") && strings.Count(line, "|") >= 1
}

func isMarkdownTableSeparator(line string) bool {
	cells := parseMarkdownTableRow(line)
	if len(cells) == 0 {
		return false
	}
	for _, cell := range cells {
		cell = strings.Trim(cell, " :-")
		if cell != "" {
			return false
		}
	}
	return true
}

func parseMarkdownTableRow(line string) []string {
	line = strings.TrimSpace(line)
	line = strings.TrimPrefix(line, "|")
	line = strings.TrimSuffix(line, "|")
	raw := strings.Split(line, "|")
	cells := make([]string, 0, len(raw))
	for _, cell := range raw {
		cells = append(cells, strings.TrimSpace(cell))
	}
	return cells
}

func padPlainLine(line string, width int) string {
	if extra := width - lipglossWidth(line); extra > 0 {
		return line + strings.Repeat(" ", extra)
	}
	return line
}

func stripANSIForWidth(line string) string {
	return stripChatANSI(line)
}

func lipglossWidth(line string) int {
	return lipgloss.Width(line)
}
