// Package markup renders markdown then sanitizes it, and never trusts the renderer. Chapter 42.
package markup

import (
	"bytes"
	"html"
	"io"
	"strconv"
	"strings"

	"github.com/yuin/goldmark"
	"github.com/yuin/goldmark/extension"
	xhtml "golang.org/x/net/html"
	"golang.org/x/net/html/atom"
)

// allowed is chapter 42.1's list, and an allowlist, so it strips elements not yet invented.
var allowed = map[atom.Atom]bool{
	atom.P: true, atom.Br: true, atom.Strong: true, atom.Em: true, atom.Del: true,
	atom.Code: true, atom.Pre: true, atom.Blockquote: true,
	atom.H1: true, atom.H2: true, atom.H3: true, atom.H4: true, atom.H5: true, atom.H6: true,
	atom.Ul: true, atom.Ol: true, atom.Li: true,
	atom.A: true, atom.Img: true,
	atom.Table: true, atom.Thead: true, atom.Tbody: true, atom.Tr: true,
	atom.Th: true, atom.Td: true,
	atom.Hr: true,
}

// allowedAttrs is the complete attribute list from chapter 42.1.
var allowedAttrs = map[atom.Atom]map[string]bool{
	atom.A:    {"href": true, "title": true},
	atom.Img:  {"src": true, "alt": true, "title": true},
	atom.Code: {"class": true}, // language-* only, checked below
	atom.Ol:   {"start": true}, // digits only, checked below
}

// dropWhole takes the contents too, or a stripped <script> leaves its source on the page.
var dropWhole = map[atom.Atom]bool{
	atom.Script: true, atom.Style: true, atom.Iframe: true, atom.Object: true,
	atom.Embed: true, atom.Form: true, atom.Svg: true, atom.Math: true,
	atom.Template: true, atom.Noscript: true,
}

// void elements have no closing tag.
var void = map[atom.Atom]bool{atom.Br: true, atom.Img: true, atom.Hr: true}

// Render turns a comment body into HTML that is safe to put on a page.
func Render(markdown string) string {
	var buf bytes.Buffer
	// GFM adds tables and strikethrough, both allowed, and raw html stays off by default.
	md := goldmark.New(goldmark.WithExtensions(extension.GFM))
	if err := md.Convert([]byte(markdown), &buf); err != nil {
		// A body that will not render is shown as the text it is, escaped.
		return "<p>" + html.EscapeString(markdown) + "</p>"
	}
	return Sanitize(buf.String())
}

// Sanitize strips everything chapter 42.1 does not allow.
func Sanitize(unsafe string) string {
	var out strings.Builder
	z := xhtml.NewTokenizer(strings.NewReader(unsafe))
	// skipped holds the tags being dropped whole, so the matching end tag is the one that resumes.
	var skipped []atom.Atom
	var open []atom.Atom

	for {
		switch z.Next() {
		case xhtml.ErrorToken:
			if z.Err() != io.EOF {
				return "" // a body that will not parse is not shown at all
			}
			// Close anything the input left open.
			for i := len(open) - 1; i >= 0; i-- {
				out.WriteString("</" + open[i].String() + ">")
			}
			return out.String()

		case xhtml.TextToken:
			if len(skipped) == 0 {
				out.WriteString(html.EscapeString(string(z.Text())))
			}

		case xhtml.StartTagToken:
			name, hasAttr := z.TagName()
			a := atom.Lookup(name)
			if len(skipped) > 0 || dropWhole[a] {
				if !void[a] {
					skipped = append(skipped, a)
				}
				continue
			}
			if !allowed[a] {
				continue // strip the tag, keep what is inside it
			}
			if writeTag(&out, z, a, hasAttr) && !void[a] {
				open = append(open, a)
			}

		case xhtml.SelfClosingTagToken:
			name, hasAttr := z.TagName()
			a := atom.Lookup(name)
			if len(skipped) > 0 || dropWhole[a] || !allowed[a] {
				continue
			}
			writeTag(&out, z, a, hasAttr)

		case xhtml.EndTagToken:
			name, _ := z.TagName()
			a := atom.Lookup(name)
			if len(skipped) > 0 {
				// Unwind to the tag this closes, so what it left open is dropped with it.
				for i := len(skipped) - 1; i >= 0; i-- {
					if skipped[i] == a {
						skipped = skipped[:i]
						break
					}
				}
				continue
			}
			if !allowed[a] || void[a] {
				continue
			}
			// Only close a tag this sanitizer actually opened.
			for i := len(open) - 1; i >= 0; i-- {
				if open[i] == a {
					out.WriteString("</" + a.String() + ">")
					open = append(open[:i], open[i+1:]...)
					break
				}
			}
		}
	}
}

// blockedImage is chapter 42.2's "say so in the UI", because a silent gap looks like a broken page.
const blockedImage = `<span class="blocked">remote image blocked, it would tell its host who read this</span>`

// writeTag emits one start tag, and reports whether it wrote one, since a blocked image writes none.
func writeTag(out *strings.Builder, z *xhtml.Tokenizer, a atom.Atom, hasAttr bool) bool {
	permitted := allowedAttrs[a]
	external, blocked := false, false
	var attrs strings.Builder
	for hasAttr {
		var k, v []byte
		k, v, hasAttr = z.TagAttr()
		key := strings.ToLower(string(k))
		value := string(v)

		// on* and style are on no allowlist, and this strips them from everything, explicitly.
		if strings.HasPrefix(key, "on") || key == "style" || !permitted[key] {
			continue
		}
		switch {
		case key == "href" || key == "src":
			clean, ok := safeURL(value)
			if !ok {
				continue
			}
			if a == atom.Img && isRemote(clean) {
				// A remote image leaks the reader's address to whoever wrote the comment. 42.2.
				blocked = true
				continue
			}
			if isRemote(clean) {
				external = true
			}
			value = clean
		case key == "start" && a == atom.Ol:
			// A list that starts at six says six, and digits are the whole of what start may hold.
			if _, err := strconv.Atoi(value); err != nil {
				continue
			}
		case key == "class" && a == atom.Code:
			if !strings.HasPrefix(value, "language-") {
				continue
			}
		}
		attrs.WriteString(" " + key + `="` + html.EscapeString(value) + `"`)
	}
	if blocked {
		out.WriteString(blockedImage)
		return false
	}
	out.WriteString("<" + a.String() + attrs.String())
	if a == atom.A && external {
		out.WriteString(` rel="nofollow noopener noreferrer"`)
	}
	out.WriteString(">")
	return true
}

// safeURL allows http, https and mailto, checked after decoding, since a browser decodes too.
func safeURL(raw string) (string, bool) {
	decoded := html.UnescapeString(raw)
	trimmed := strings.TrimLeft(decoded, " \t\r\n\x00\v\f")
	lower := strings.ToLower(trimmed)

	// A scheme is what precedes the first colon, if it comes before any slash, question or hash.
	if i := strings.IndexAny(lower, ":/?#"); i >= 0 && lower[i] == ':' {
		scheme := lower[:i]
		switch scheme {
		case "http", "https", "mailto":
		default:
			return "", false
		}
	}
	if strings.ContainsAny(trimmed, "\x00\n\r") {
		return "", false
	}
	return trimmed, true
}

func isRemote(u string) bool {
	l := strings.ToLower(u)
	return strings.HasPrefix(l, "http://") || strings.HasPrefix(l, "https://") ||
		strings.HasPrefix(l, "//")
}