// Package htmlmd turns a web page into Markdown for the model to read: its
// headings, paragraphs, lists, tables and code, without scripts, styles,
// navigation, images or where links lead. The page is untrusted; parsing
// it is x/net/html's, which follows the HTML5 spec, and this package only
// walks the tree that makes.
package htmlmd

import (
	"fmt"
	"io"
	"strings"
	"unicode"

	"golang.org/x/net/html"
	"golang.org/x/net/html/atom"
)

// Page is a converted page.
type Page struct {
	Markdown string
	Headings []Heading
	// Anchors are the lines (from 1) where the elements with each id start,
	// for finding a URL's #fragment.
	Anchors map[string]int
}

// Heading is one of a page's headings.
type Heading struct {
	Line  int // from 1
	Level int // 1 to 6
	Text  string
}

// Convert reads an HTML page and turns it into Markdown.
func Convert(r io.Reader) (Page, error) {
	doc, err := html.Parse(r)
	if err != nil {
		return Page{}, err
	}
	w := &writer{page: Page{Anchors: map[string]int{}}}
	w.walk(doc)
	w.flush()
	w.page.Markdown = w.out.String()
	return w.page, nil
}

type writer struct {
	page  Page
	out   strings.Builder
	lines int
	// text gathers the words of the block being read, until a block
	// boundary writes them out as a line.
	text strings.Builder
	// prefix is what begins each line in the open lists and quotes. A list
	// item's first line has its marker in place of its indent.
	prefix []string
	marker string
	// lists holds, for each open list, the number of its next item, or 0
	// for a bulleted list.
	lists []int
	// blank asks for a blank line before the next block, outside lists.
	blank bool
	// pending are ids whose elements start on the next line written.
	pending []string
}

// skipped are elements with nothing to read in them.
var skipped = map[atom.Atom]bool{
	atom.Head: true, atom.Script: true, atom.Style: true, atom.Noscript: true,
	atom.Template: true, atom.Nav: true, atom.Footer: true, atom.Svg: true,
	atom.Math: true, atom.Iframe: true, atom.Object: true, atom.Embed: true,
	atom.Img: true, atom.Picture: true, atom.Video: true, atom.Audio: true,
	atom.Canvas: true, atom.Form: true, atom.Button: true, atom.Input: true,
	atom.Select: true, atom.Textarea: true, atom.Dialog: true,
}

// skippedRoles are the ARIA roles of what nav, footer and a search form
// hold, for pages that mark them with a role instead.
var skippedRoles = map[string]bool{"navigation": true, "search": true, "contentinfo": true}

// blocks are elements that start and end a block of text.
var blocks = map[atom.Atom]bool{
	atom.P: true, atom.Div: true, atom.Section: true, atom.Article: true,
	atom.Main: true, atom.Aside: true, atom.Header: true, atom.Figure: true,
	atom.Figcaption: true, atom.Details: true, atom.Summary: true,
	atom.Address: true, atom.Center: true, atom.Body: true, atom.Html: true,
	atom.Caption: true, atom.Fieldset: true, atom.Legend: true,
	atom.Dt: true, atom.Hgroup: true,
}

// skip reports whether an element is left out, with all it holds: skipped,
// hidden, or in a skipped role.
func skip(n *html.Node) bool {
	return skipped[n.DataAtom] || hasAttr(n, "hidden") || attr(n, "aria-hidden") == "true" || skippedRoles[attr(n, "role")]
}

var headings = map[atom.Atom]int{
	atom.H1: 1, atom.H2: 2, atom.H3: 3, atom.H4: 4, atom.H5: 5, atom.H6: 6,
}

func (w *writer) walk(n *html.Node) {
	switch n.Type {
	case html.DocumentNode:
		w.children(n)
		return
	case html.TextNode:
		w.text.WriteString(n.Data)
		return
	case html.ElementNode:
	default: // comments, doctypes
		return
	}
	if skip(n) {
		return
	}
	inline := !blocks[n.DataAtom] && headings[n.DataAtom] == 0
	switch n.DataAtom {
	case atom.Pre, atom.Ul, atom.Ol, atom.Li, atom.Dl, atom.Dd, atom.Blockquote, atom.Table, atom.Hr, atom.Br:
		inline = false
	}
	if !inline {
		w.flush()
	}
	// An inline element's id is on the line its block will be written to.
	if id := attr(n, "id"); id != "" {
		w.pending = append(w.pending, id)
	}
	if level := headings[n.DataAtom]; level > 0 {
		w.heading(n, level)
		return
	}
	switch n.DataAtom {
	case atom.Br:
		// A line break ends a line, but not a block.
		w.blank = false
	case atom.Hr:
		w.separate()
		w.line("---")
		w.blank = true
	case atom.Pre:
		w.pre(n)
	case atom.Table:
		w.table(n)
	case atom.Code, atom.Kbd, atom.Samp, atom.Tt:
		w.text.WriteString("`")
		w.children(n)
		w.text.WriteString("`")
	case atom.Ul, atom.Ol:
		w.separate()
		next := 0
		if n.DataAtom == atom.Ol {
			next = 1
		}
		w.lists = append(w.lists, next)
		w.children(n)
		w.flush()
		w.lists = w.lists[:len(w.lists)-1]
		w.blank = len(w.lists) == 0
	case atom.Li:
		marker := "- "
		if k := len(w.lists) - 1; k >= 0 && w.lists[k] > 0 {
			marker = fmt.Sprintf("%d. ", w.lists[k])
			w.lists[k]++
		}
		w.indented(n, marker)
	case atom.Dd:
		w.indented(n, "")
	case atom.Blockquote:
		w.separate()
		w.prefix = append(w.prefix, "> ")
		w.children(n)
		w.flush()
		w.prefix = w.prefix[:len(w.prefix)-1]
		w.blank = true
	default:
		w.children(n)
		if !inline {
			w.flush()
		}
	}
}

func (w *writer) children(n *html.Node) {
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		w.walk(c)
	}
}

// indented writes a list item, or a definition, indented under the line
// before it; marker begins its first line.
func (w *writer) indented(n *html.Node, marker string) {
	indent := strings.Repeat(" ", max(len(marker), 2))
	w.prefix = append(w.prefix, indent)
	w.marker = marker
	w.children(n)
	w.flush()
	w.prefix = w.prefix[:len(w.prefix)-1]
	w.marker = ""
}

// heading writes a heading, with the ids inside it, since they name it.
func (w *writer) heading(n *html.Node, level int) {
	w.ids(n)
	text := words(textOf(n, false))
	if text == "" {
		return
	}
	w.separate()
	w.line(strings.Repeat("#", level) + " " + text)
	w.page.Headings = append(w.page.Headings, Heading{Line: w.lines, Level: level, Text: text})
	w.blank = true
}

// pre writes preformatted text as a fenced code block, as it is.
func (w *writer) pre(n *html.Node) {
	w.ids(n)
	text := strings.Trim(textOf(n, true), "\n")
	if strings.TrimSpace(text) == "" {
		return
	}
	// The fence must be longer than any run of backticks in the code.
	fence := "```"
	for strings.Contains(text, fence) {
		fence += "`"
	}
	w.separate()
	w.line(fence)
	for l := range strings.SplitSeq(text, "\n") {
		w.line(Clean(strings.TrimRight(l, " \t\r")))
	}
	w.line(fence)
	w.blank = true
}

// table writes a table's rows, the first one as its header. A table in a
// cell is read as the cell's text.
func (w *writer) table(n *html.Node) {
	w.ids(n)
	var rows [][]string
	var collect func(*html.Node)
	collect = func(n *html.Node) {
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			if c.Type != html.ElementNode || skip(c) {
				continue
			}
			if c.DataAtom != atom.Tr {
				collect(c) // thead, tbody, tfoot
				continue
			}
			var row []string
			for cell := c.FirstChild; cell != nil; cell = cell.NextSibling {
				if cell.Type == html.ElementNode && (cell.DataAtom == atom.Td || cell.DataAtom == atom.Th) {
					row = append(row, strings.ReplaceAll(words(textOf(cell, false)), "|", `\|`))
				}
			}
			if len(row) > 0 {
				rows = append(rows, row)
			}
		}
	}
	collect(n)
	if len(rows) == 0 {
		return
	}
	w.separate()
	for i, row := range rows {
		w.line("| " + strings.Join(row, " | ") + " |")
		if i == 0 {
			w.line(strings.Repeat("| --- ", len(row)) + "|")
		}
	}
	w.blank = true
}

// ids marks the ids in a subtree that is written as a whole as starting
// on its first line.
func (w *writer) ids(n *html.Node) {
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		if c.Type == html.ElementNode {
			if id := attr(c, "id"); id != "" {
				w.pending = append(w.pending, id)
			}
			w.ids(c)
		}
	}
}

// flush writes the gathered words as a line.
func (w *writer) flush() {
	text := words(w.text.String())
	w.text.Reset()
	if text == "" {
		return
	}
	w.separate()
	w.line(text)
	w.blank = true
}

// separate writes the blank line a block asked for, except inside lists,
// whose items stay together.
func (w *writer) separate() {
	if w.blank && w.lines > 0 && len(w.lists) == 0 {
		w.out.WriteString(strings.TrimRight(strings.Join(w.prefix, ""), " "))
		w.out.WriteByte('\n')
		w.lines++
	}
	w.blank = false
}

// line writes a line after the prefix of the lists and quotes it is in.
func (w *writer) line(s string) {
	p := strings.Join(w.prefix, "")
	if w.marker != "" && strings.HasSuffix(p, strings.Repeat(" ", len(w.marker))) {
		p = p[:len(p)-len(w.marker)] + w.marker
	}
	w.marker = ""
	w.out.WriteString(strings.TrimRight(p+s, " "))
	w.out.WriteByte('\n')
	w.lines++
	for _, id := range w.pending {
		if _, ok := w.page.Anchors[id]; !ok {
			w.page.Anchors[id] = w.lines
		}
	}
	w.pending = w.pending[:0]
}

// textOf is the text in a subtree, leaving out what skip does. pre keeps
// line breaks (<br> is one); otherwise they are spaces, for words to fold.
func textOf(n *html.Node, pre bool) string {
	var b strings.Builder
	var walk func(*html.Node)
	walk = func(n *html.Node) {
		switch {
		case n.Type == html.TextNode:
			b.WriteString(n.Data)
		case n.Type != html.ElementNode || skip(n):
		case n.DataAtom == atom.Br:
			if pre {
				b.WriteString("\n")
			} else {
				b.WriteString(" ")
			}
		default:
			for c := n.FirstChild; c != nil; c = c.NextSibling {
				walk(c)
			}
		}
	}
	walk(n)
	return b.String()
}

// words folds runs of space, line breaks included, into single spaces.
// It drops the pilcrows sites link headings and definitions to themselves
// with, and control characters, which no page needs and a terminal might
// act on.
func words(s string) string {
	return strings.Join(strings.Fields(strings.ReplaceAll(Clean(s), "¶", " ")), " ")
}

// Clean drops control characters but line breaks and tabs.
func Clean(s string) string {
	return strings.Map(func(r rune) rune {
		if unicode.IsControl(r) && r != '\n' && r != '\t' {
			return -1
		}
		return r
	}, s)
}

func attr(n *html.Node, key string) string {
	for _, a := range n.Attr {
		if a.Namespace == "" && a.Key == key {
			return a.Val
		}
	}
	return ""
}

func hasAttr(n *html.Node, key string) bool {
	for _, a := range n.Attr {
		if a.Namespace == "" && a.Key == key {
			return true
		}
	}
	return false
}