package htmlformat // import alin.ovh/htmlformat

import (
	"bufio"
	"fmt"
	"io"
	"strings"
	"unicode"
	"unicode/utf8"

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

// IndentString is the string used for indentation. Use "\t" for tab.
var IndentString = "  "

// Document formats a HTML document.
func Document(w io.Writer, r io.Reader) (err error) {
	node, err := html.Parse(r)
	if err != nil {
		return err
	}
	return Nodes(w, []*html.Node{node})
}

// Fragment formats a fragment of a HTML document.
func Fragment(w io.Writer, r io.Reader) (err error) {
	context := &html.Node{
		Type: html.ElementNode,
	}
	nodes, err := html.ParseFragment(r, context)
	if err != nil {
		return err
	}
	return Nodes(w, nodes)
}

// Nodes formats a slice of HTML nodes.
func Nodes(w io.Writer, nodes []*html.Node) (err error) {
	for _, node := range nodes {
		if err = printNode(w, node, false, 0); err != nil {
			return
		}
	}
	return
}

// Is this node a tag with no end tag such as <meta> or <br>?
// http://www.w3.org/TR/html-markup/syntax.html#syntax-elements
func isVoidElement(n *html.Node) bool {
	switch n.DataAtom {
	case atom.Area, atom.Base, atom.Br, atom.Col, atom.Command, atom.Embed,
		atom.Hr, atom.Img, atom.Input, atom.Keygen, atom.Link,
		atom.Meta, atom.Param, atom.Source, atom.Track, atom.Wbr:
		return true
	}
	return false
}

func isInlineElement(n *html.Node) bool {
	switch n.DataAtom {
	case atom.B, atom.I, atom.U, atom.S,
		atom.A, atom.Br, atom.Code, atom.Em, atom.Time,
		atom.Span, atom.Strong, atom.Small, atom.Mark, atom.Del,
		atom.Ins, atom.Sub, atom.Sup, atom.Q, atom.Cite,
		atom.Dfn, atom.Abbr, atom.Data, atom.Var, atom.Samp,
		atom.Kbd, atom.Label, atom.Button, atom.Select, atom.Textarea,
		atom.Img, atom.Map, atom.Object, atom.Iframe, atom.Audio,
		atom.Video, atom.Canvas, atom.Meter, atom.Progress, atom.Math:
		return true
	}
	return false
}
func isSpecialContentElement(n *html.Node) bool {
	if n != nil {
		switch n.DataAtom {
		case atom.Style,
			atom.Script:
			return true
		}
	}
	return false
}

func collapseWhitespace(in string) string {
	if in == "" {
		return ""
	}

	leading := unicode.IsSpace(getFirstRune(in))
	trailing := unicode.IsSpace(getLastRune(in))

	out := strings.TrimSpace(in)
	switch {
	case leading && trailing:
		return " " + out + " "
	case leading:
		return " " + out
	case trailing:
		return out + " "
	default:
		return out
	}
}

func getFirstRune(s string) rune {
	r, _ := utf8.DecodeRuneInString(s)
	return r
}

func getLastRune(s string) rune {
	r, _ := utf8.DecodeLastRuneInString(s)
	return r
}

func hasSingleTextChild(n *html.Node) bool {
	return n != nil && n.FirstChild != nil && n.FirstChild == n.LastChild &&
		n.FirstChild.Type == html.TextNode
}

func printNode(w io.Writer, n *html.Node, pre bool, level int) (err error) {
	switch n.Type {
	case html.TextNode:
		if pre {
			if _, err = fmt.Fprint(w, html.EscapeString(n.Data)); err != nil {
				return
			}
			return nil
		}
		s := n.Data
		s = strings.TrimSpace(s)
		if s != "" {
			if !isSpecialContentElement(n.Parent) && !hasSingleTextChild(n.Parent) &&
				(n.PrevSibling == nil || !unicode.IsPunct(getFirstRune(s))) {
				if err = printIndent(w, level); err != nil {
					return
				}
			}
			if isSpecialContentElement(n.Parent) {
				scanner := bufio.NewScanner(strings.NewReader(s))
				for scanner.Scan() {
					t := scanner.Text()
					if _, err = fmt.Fprintln(w); err != nil {
						return
					}
					if err = printIndent(w, level); err != nil {
						return
					}
					if _, err = fmt.Fprint(w, t); err != nil {
						return
					}
				}
				if err = scanner.Err(); err != nil {
					return
				}
				if _, err = fmt.Fprintln(w); err != nil {
					return
				}
			} else {
				if _, err = fmt.Fprint(w, collapseWhitespace(s)); err != nil {
					return
				}
				if !hasSingleTextChild(n.Parent) {
					if unicode.IsSpace(getLastRune(n.Data)) || n.NextSibling == nil || (n.NextSibling.Type == html.ElementNode && !isInlineElement(n.NextSibling)) {
						if _, err = fmt.Fprint(w, "\n"); err != nil {
							return
						}
					}
				}
			}
		}
	case html.ElementNode:
		if !pre && (n.PrevSibling == nil ||
			(n.PrevSibling.Type != html.TextNode || unicode.IsSpace(getLastRune(n.PrevSibling.Data)))) {
			if err = printIndent(w, level); err != nil {
				return
			}
		}
		if _, err = fmt.Fprintf(w, "<%s", n.Data); err != nil {
			return
		}
		for _, a := range n.Attr {
			val := html.EscapeString(a.Val)
			if _, err = fmt.Fprintf(w, ` %s="%s"`, a.Key, val); err != nil {
				return
			}
		}
		if _, err = fmt.Fprint(w, ">"); err != nil {
			return
		}
		if !pre && !hasSingleTextChild(n) {
			if _, err = fmt.Fprint(w, "\n"); err != nil {
				return
			}
		}
		if !isVoidElement(n) {
			if err = printChildren(w, n, pre || n.Data == "pre", level+1); err != nil {
				return
			}
			if !pre && (isSpecialContentElement(n) || !hasSingleTextChild(n)) {
				if err = printIndent(w, level); err != nil {
					return
				}
			}
			if _, err = fmt.Fprintf(w, "</%s>", n.Data); err != nil {
				return
			}

			if !pre && (n.NextSibling == nil ||
				(n.NextSibling.Type == html.ElementNode) ||
				(n.NextSibling.Type == html.TextNode && !unicode.IsPunct(getFirstRune(n.NextSibling.Data)))) {
				if _, err = fmt.Fprint(w, "\n"); err != nil {
					return
				}
			}
		}
	case html.CommentNode:
		if err = printIndent(w, level); err != nil {
			return
		}
		if _, err = fmt.Fprintf(w, "<!--%s-->\n", n.Data); err != nil {
			return
		}
		if err = printChildren(w, n, false, level); err != nil {
			return
		}
	case html.DoctypeNode:
		if _, err = fmt.Fprintf(w, "<!doctype %s>\n", n.Data); err != nil {
			return
		}
		if err = printChildren(w, n, false, level); err != nil {
			return
		}
	case html.DocumentNode:
		if err = printChildren(w, n, false, level); err != nil {
			return
		}
	}
	return
}

func printChildren(w io.Writer, n *html.Node, pre bool, level int) (err error) {
	child := n.FirstChild
	for child != nil {
		if err = printNode(w, child, pre, level); err != nil {
			return
		}
		child = child.NextSibling
	}
	return
}

func printIndent(w io.Writer, level int) (err error) {
	_, err = fmt.Fprint(w, strings.Repeat(IndentString, level))
	return err
}