package classes

import (
	"bytes"
	"sort"
	"strings"

	"github.com/yuin/goldmark/ast"
	gmtext "github.com/yuin/goldmark/text"
	"github.com/yuin/goldmark/util"
	xhtml "golang.org/x/net/html"
)

// SetImageAlt sets the alt attribute of every img in html whose src
// attribute (unescaped) equals src, adding it when missing. It returns the
// new HTML, the number of matching images and the alt the first match had
// (nil when it had none). It works on the x/net/html tokenizer: a matching
// img tag is written again from its token, so every attribute value,
// alt included, is escaped by the tokenizer and never concatenated; every
// other byte of html is copied as it is, the bytes of an unfinished tag at
// the end included. An empty alt marks a decorative image (alt="").
func SetImageAlt(html, src, alt string) (string, int, *string) {
	if src == "" || !containsFold(html, "<img") {
		return html, 0, nil
	}
	z := xhtml.NewTokenizer(strings.NewReader(html))
	var b bytes.Buffer
	b.Grow(len(html) + len(alt))
	matches := 0
	var old *string
	consumed := 0
	for {
		tt := z.Next()
		if tt == xhtml.ErrorToken {
			// The bytes of an unfinished tag at the end are kept.
			b.WriteString(html[min(consumed, len(html)):])
			break
		}
		raw := z.Raw()
		consumed += len(raw)
		if tt == xhtml.StartTagToken || tt == xhtml.SelfClosingTagToken {
			// Raw is copied before Token, which reads the tag; Token is
			// the only reader of the tag (TagName would consume it).
			kept := append([]byte(nil), raw...)
			tok := z.Token()
			if tok.Data == "img" && imgSrc(tok) == src {
				prev := setAlt(&tok, alt)
				if matches == 0 {
					old = prev
				}
				matches++
				b.WriteString(tok.String())
				continue
			}
			b.Write(kept)
			continue
		}
		b.Write(raw)
	}
	if matches == 0 {
		return html, 0, nil
	}
	return b.String(), matches, old
}

func imgSrc(tok xhtml.Token) string {
	for _, a := range tok.Attr {
		if a.Namespace == "" && a.Key == "src" {
			return a.Val
		}
	}
	return ""
}

// setAlt sets tok's alt to alt and returns the previous value (nil when
// there was none). Repeated alt attributes collapse to one.
func setAlt(tok *xhtml.Token, alt string) *string {
	var prev *string
	attrs := tok.Attr[:0]
	set := false
	for _, a := range tok.Attr {
		if a.Namespace == "" && a.Key == "alt" {
			if prev == nil {
				v := a.Val
				prev = &v
			}
			if set {
				continue
			}
			a.Val = alt
			set = true
		}
		attrs = append(attrs, a)
	}
	if !set {
		attrs = append(attrs, xhtml.Attribute{Key: "alt", Val: alt})
	}
	tok.Attr = attrs
	return prev
}

// SetMarkdownImageAlt sets the alt text of every inline markdown image of
// md whose destination is src (P33-65): ![label](dest) and
// ![label](dest "title"), outside fenced and indented code blocks, code
// spans and HTML blocks (the ranges come from goldmark's AST, the engine of
// FormatHTML). An image matches when its destination equals src as written,
// or as goldmark's HTML renderer escapes it into the img src attribute. The
// label is replaced by alt with \, [ and ] backslash-escaped and line breaks
// folded to spaces, so an alt can neither close the image nor add a link;
// every other byte of md is kept. It returns the new markdown, the number of
// matching images and the label the first match had. A reference-style
// image (![label][ref]) is not matched.
func SetMarkdownImageAlt(md, src, alt string) (string, int, *string) {
	if src == "" || !strings.Contains(md, "![") {
		return md, 0, nil
	}
	skip, images := markdownLayout(md)
	if len(images) == 0 {
		return md, 0, nil
	}
	label := escapeMarkdownLabel(alt)
	var b strings.Builder
	b.Grow(len(md) + len(label))
	matches := 0
	var old *string
	last := 0
	for i := 0; i+1 < len(md); i++ {
		if md[i] != '!' || md[i+1] != '[' || (i > 0 && md[i-1] == '\\') || inRanges(skip, i) {
			continue
		}
		img, ok := scanInlineImage(md, i)
		if !ok {
			continue
		}
		dest := md[img.destStart:img.destEnd]
		escaped := string(util.URLEscape([]byte(dest), true))
		if !images[dest] || (dest != src && escaped != src) {
			i = img.end - 1
			continue
		}
		if matches == 0 {
			v := md[img.labelStart:img.labelEnd]
			old = &v
		}
		matches++
		b.WriteString(md[last:img.labelStart])
		b.WriteString(label)
		last = img.labelEnd
		i = img.end - 1
	}
	if matches == 0 {
		return md, 0, nil
	}
	b.WriteString(md[last:])
	return b.String(), matches, old
}

// SetImageAltOutsideCode applies SetImageAlt to the parts of markdown md
// outside its fenced and indented code blocks and code spans, so a raw
// <img> in the markdown gets the alt while an <img> shown in a code sample
// keeps its bytes.
func SetImageAltOutsideCode(md, src, alt string) (string, int, *string) {
	if src == "" || !containsFold(md, "<img") {
		return md, 0, nil
	}
	code := markdownCodeRanges(md)
	var b strings.Builder
	b.Grow(len(md) + len(alt))
	matches := 0
	var old *string
	last := 0
	apply := func(part string) {
		out, n, prev := SetImageAlt(part, src, alt)
		if n > 0 && matches == 0 {
			old = prev
		}
		matches += n
		b.WriteString(out)
	}
	for _, r := range code {
		apply(md[last:r[0]])
		b.WriteString(md[r[0]:r[1]])
		last = r[1]
	}
	apply(md[last:])
	if matches == 0 {
		return md, 0, nil
	}
	return b.String(), matches, old
}

// escapeMarkdownLabel is alt as an image label: \, [ and ] escaped, line
// breaks folded to spaces.
func escapeMarkdownLabel(alt string) string {
	alt = strings.NewReplacer("\r\n", " ", "\n", " ", "\r", " ").Replace(alt)
	return strings.NewReplacer(`\`, `\\`, "[", `\[`, "]", `\]`).Replace(alt)
}

// inlineImage is the byte layout of one ![label](dest "title") in a text.
type inlineImage struct {
	labelStart, labelEnd int
	destStart, destEnd   int
	end                  int
}

// scanInlineImage reads the inline image starting at md[i] ("!["), as
// goldmark's link parser reads one: a label with balanced brackets and
// backslash escapes, "(", optional spaces, a destination in <...> or
// without spaces with balanced parentheses, an optional title in quotes or
// parentheses, and ")".
func scanInlineImage(md string, i int) (inlineImage, bool) {
	var img inlineImage
	img.labelStart = i + 2
	depth := 0
	j := img.labelStart
	for ; j < len(md); j++ {
		c := md[j]
		if c == '\\' && j+1 < len(md) {
			j++
			continue
		}
		if c == '[' {
			depth++
		} else if c == ']' {
			if depth == 0 {
				break
			}
			depth--
		}
	}
	if j >= len(md) || j+1 >= len(md) || md[j+1] != '(' {
		return img, false
	}
	img.labelEnd = j
	k := skipLinkSpace(md, j+2)
	if k >= len(md) {
		return img, false
	}
	if md[k] == '<' {
		e := k + 1
		for ; e < len(md) && md[e] != '>' && md[e] != '\n'; e++ {
			if md[e] == '\\' && e+1 < len(md) {
				e++
			}
		}
		if e >= len(md) || md[e] != '>' {
			return img, false
		}
		img.destStart, img.destEnd = k+1, e
		k = e + 1
	} else {
		opened := 0
		e := k
		for ; e < len(md); e++ {
			c := md[e]
			if c == '\\' && e+1 < len(md) {
				e++
				continue
			}
			if c == '(' {
				opened++
			} else if c == ')' {
				if opened == 0 {
					break
				}
				opened--
			} else if c == ' ' || c == '\t' || c == '\n' || c == '\r' {
				break
			}
		}
		if e == k {
			return img, false
		}
		img.destStart, img.destEnd = k, e
		k = e
	}
	k = skipLinkSpace(md, k)
	if k < len(md) && (md[k] == '"' || md[k] == '\'' || md[k] == '(') {
		closer := md[k]
		if closer == '(' {
			closer = ')'
		}
		e := k + 1
		for ; e < len(md) && md[e] != closer; e++ {
			if md[e] == '\\' && e+1 < len(md) {
				e++
			}
		}
		if e >= len(md) {
			return img, false
		}
		k = skipLinkSpace(md, e+1)
	}
	if k >= len(md) || md[k] != ')' {
		return img, false
	}
	img.end = k + 1
	return img, true
}

// skipLinkSpace skips spaces, tabs and at most one line break.
func skipLinkSpace(md string, k int) int {
	newline := false
	for k < len(md) {
		switch md[k] {
		case ' ', '\t':
		case '\n', '\r':
			if newline {
				return k
			}
			newline = true
		default:
			return k
		}
		k++
	}
	return k
}

// markdownLayout parses md as FormatHTML does and answers the byte ranges
// an image is never read from (code blocks, code spans and HTML blocks)
// and the destinations of its inline images.
func markdownLayout(md string) ([][2]int, map[string]bool) {
	images := map[string]bool{}
	ranges := markdownRanges(md, true, func(n ast.Node) {
		if img, ok := n.(*ast.Image); ok {
			images[string(img.Destination)] = true
		}
	})
	return ranges, images
}

// markdownCodeRanges is the byte ranges of md's code blocks and code
// spans, in order.
func markdownCodeRanges(md string) [][2]int {
	return markdownRanges(md, false, nil)
}

// markdownRanges walks md's goldmark AST (parsed from the trimmed text, as
// FormatHTML converts it) and answers the sorted, merged byte ranges of its
// fenced and indented code blocks, its code spans and, with html, its HTML
// blocks; visit sees every node.
func markdownRanges(md string, html bool, visit func(ast.Node)) [][2]int {
	trimmed := strings.TrimSpace(md)
	lead := strings.Index(md, trimmed)
	if trimmed == "" || lead < 0 {
		return nil
	}
	source := []byte(trimmed)
	doc := markdownEngine.Parser().Parse(gmtext.NewReader(source))
	var ranges [][2]int
	lines := func(n ast.Node) {
		l := n.Lines()
		for i := 0; i < l.Len(); i++ {
			s := l.At(i)
			ranges = append(ranges, [2]int{lead + s.Start, lead + s.Stop})
		}
	}
	_ = ast.Walk(doc, func(n ast.Node, entering bool) (ast.WalkStatus, error) {
		if !entering {
			return ast.WalkContinue, nil
		}
		if visit != nil {
			visit(n)
		}
		switch v := n.(type) {
		case *ast.FencedCodeBlock:
			lines(v)
			if v.Info != nil {
				ranges = append(ranges, [2]int{lead + v.Info.Segment.Start, lead + v.Info.Segment.Stop})
			}
			return ast.WalkSkipChildren, nil
		case *ast.CodeBlock:
			lines(v)
			return ast.WalkSkipChildren, nil
		case *ast.HTMLBlock:
			if html {
				lines(v)
			}
			return ast.WalkSkipChildren, nil
		case *ast.CodeSpan:
			start, stop := -1, -1
			for c := v.FirstChild(); c != nil; c = c.NextSibling() {
				if t, ok := c.(*ast.Text); ok {
					if start < 0 || t.Segment.Start < start {
						start = t.Segment.Start
					}
					stop = max(stop, t.Segment.Stop)
				}
			}
			if start >= 0 {
				// The backticks around the span belong to it.
				for start > 0 && source[start-1] == '`' {
					start--
				}
				for stop < len(source) && source[stop] == '`' {
					stop++
				}
				ranges = append(ranges, [2]int{lead + start, lead + stop})
			}
			return ast.WalkSkipChildren, nil
		}
		return ast.WalkContinue, nil
	})
	sort.Slice(ranges, func(i, j int) bool { return ranges[i][0] < ranges[j][0] })
	var merged [][2]int
	for _, r := range ranges {
		if r[0] >= r[1] {
			continue
		}
		if n := len(merged); n > 0 && r[0] <= merged[n-1][1] {
			merged[n-1][1] = max(merged[n-1][1], r[1])
			continue
		}
		merged = append(merged, r)
	}
	return merged
}

// inRanges reports whether i lies in one of the sorted ranges.
func inRanges(ranges [][2]int, i int) bool {
	for _, r := range ranges {
		if i < r[0] {
			return false
		}
		if i < r[1] {
			return true
		}
	}
	return false
}