package format

import (
	"strings"
	"unicode/utf8"
)

const defaultMaxLineWidth = 120

// Block elements get their own line and increase indentation for children.
var blockTags = map[string]bool{
	// Layout
	"ac:layout":         true,
	"ac:layout-section": true,
	"ac:layout-cell":    true,
	// Block content
	"p":  true,
	"h1": true, "h2": true, "h3": true, "h4": true, "h5": true, "h6": true,
	"div": true,
	// Lists
	"ul": true, "ol": true, "li": true,
	// Tables
	"table": true, "thead": true, "tbody": true, "colgroup": true,
	"tr": true, "th": true, "td": true,
	// Macros
	"ac:structured-macro": true,
	"ac:rich-text-body":   true,
	"ac:plain-text-body":  true,
	// Task lists
	"ac:task-list": true,
	"ac:task":      true,
	"ac:task-body": true,
}

// inlineableBlocks: block tags that prefer to stay on one line if short enough.
var inlineableBlocks = map[string]bool{
	"li": true, "th": true, "td": true,
	"h1": true, "h2": true, "h3": true, "h4": true, "h5": true, "h6": true,
	"ac:task-id": true, "ac:task-status": true,
}

// Pre elements: content inside is not reformatted.
var preTags = map[string]bool{
	"ac:plain-text-body": true,
}

// PrettyXML formats Confluence storage XML with sensible indentation.
func PrettyXML(input string, indent string) string {
	tokens := tokenize(input)
	var buf strings.Builder
	level := 0
	inPre := 0
	atLineStart := true

	i := 0
	for i < len(tokens) {
		tok := tokens[i]

		switch tok.kind {
		case tokenOpen:
			tagName := tok.tagName()
			if inPre > 0 {
				buf.WriteString(tok.raw)
				if preTags[tagName] {
					inPre++
				}
				i++
				continue
			}
			if preTags[tagName] {
				inPre++
				ensureIndentedLine(&buf, level, indent, &atLineStart)
				buf.WriteString(tok.raw)
				i++
				continue
			}
			if blockTags[tagName] {
				// Try to inline short blocks like <li>text</li>, <h1>Title</h1>
				if inlineableBlocks[tagName] {
					if inlined, skip := tryInlineBlock(tokens[i:], tagName); skip > 0 {
						ensureIndentedLine(&buf, level, indent, &atLineStart)
						buf.WriteString(inlined)
						buf.WriteString("\n")
						atLineStart = true
						i += skip
						continue
					}
				}
				ensureIndentedLine(&buf, level, indent, &atLineStart)
				buf.WriteString(tok.raw)
				buf.WriteString("\n")
				level++
				atLineStart = true
			} else {
				if atLineStart {
					writeIndentPrefix(&buf, level, indent)
					atLineStart = false
				}
				buf.WriteString(tok.raw)
			}

		case tokenClose:
			tagName := tok.tagName()
			if inPre > 0 {
				buf.WriteString(tok.raw)
				if preTags[tagName] {
					inPre--
				}
				i++
				continue
			}
			if blockTags[tagName] {
				level--
				if level < 0 {
					level = 0
				}
				if !atLineStart {
					buf.WriteString("\n")
				}
				writeIndentPrefix(&buf, level, indent)
				buf.WriteString(tok.raw)
				buf.WriteString("\n")
				atLineStart = true
			} else {
				buf.WriteString(tok.raw)
			}

		case tokenSelfClose:
			tagName := tok.tagName()
			raw := normalizeSelfClose(tok.raw)
			if inPre > 0 {
				buf.WriteString(raw)
				i++
				continue
			}
			if blockTags[tagName] || tagName == "hr" || tagName == "col" {
				ensureIndentedLine(&buf, level, indent, &atLineStart)
				buf.WriteString(raw)
				buf.WriteString("\n")
				atLineStart = true
			} else {
				if atLineStart {
					writeIndentPrefix(&buf, level, indent)
					atLineStart = false
				}
				buf.WriteString(raw)
			}

		case tokenText:
			if inPre > 0 {
				buf.WriteString(tok.raw)
				i++
				continue
			}
			text := collapseWS(tok.raw)
			if text == "" {
				i++
				continue
			}
			// A lone space is an indentation artifact only at a block boundary
			// (line start); mid-line it sits between two inline elements
			// (e.g. </strong> <code>) and is a significant space — keep it.
			if atLineStart {
				text = strings.TrimLeft(text, " ")
				if text == "" {
					i++
					continue
				}
				writeIndentPrefix(&buf, level, indent)
				atLineStart = false
			}
			buf.WriteString(text)

		case tokenCDATA, tokenComment:
			if inPre > 0 {
				buf.WriteString(tok.raw)
				i++
				continue
			}
			if atLineStart {
				writeIndentPrefix(&buf, level, indent)
				atLineStart = false
			}
			buf.WriteString(tok.raw)
		}

		i++
	}

	result := normalizeEntities(buf.String())
	// Post-process: clean up lines and wrap long ones
	lines := strings.Split(result, "\n")
	var final []string
	for _, line := range lines {
		line = strings.TrimRight(line, " \t")
		if runeWidth(line) > defaultMaxLineWidth {
			final = append(final, wrapLine(line, defaultMaxLineWidth)...)
		} else {
			final = append(final, line)
		}
	}
	return strings.TrimSpace(strings.Join(final, "\n")) + "\n"
}

// tryInlineBlock checks if the block starting at tokens[0] (an open tag) has
// only inline/text children and a matching close tag, and the total is short
// enough to fit on one line. Returns the inlined string and number of tokens consumed.
func tryInlineBlock(tokens []token, tagName string) (string, int) {
	if len(tokens) < 2 {
		return "", 0
	}
	// Scan forward to find matching close tag
	depth := 0
	var inner strings.Builder
	for j, tok := range tokens {
		if j == 0 {
			inner.WriteString(tok.raw)
			depth = 1
			continue
		}
		switch tok.kind {
		case tokenOpen:
			tn := tok.tagName()
			if blockTags[tn] && !inlineableBlocks[tn] {
				// Contains a non-inlineable block child — can't inline
				return "", 0
			}
			if tn == tagName {
				depth++
			}
			inner.WriteString(tok.raw)
		case tokenClose:
			tn := tok.tagName()
			if tn == tagName {
				depth--
				if depth == 0 {
					inner.WriteString(tok.raw)
					result := inner.String()
					if runeWidth(result) <= defaultMaxLineWidth {
						return result, j + 1
					}
					return "", 0
				}
			}
			inner.WriteString(tok.raw)
		case tokenText:
			text := collapseWS(tok.raw)
			if text == "" {
				continue
			}
			// Trim leading space only for the first text token after open tag
			if j == 1 {
				text = strings.TrimLeft(text, " ")
			}
			inner.WriteString(text)
		case tokenCDATA:
			// CDATA in an inlineable block — don't inline if multiline
			if strings.Contains(tok.raw, "\n") {
				return "", 0
			}
			inner.WriteString(tok.raw)
		case tokenSelfClose:
			inner.WriteString(normalizeSelfClose(tok.raw))
		default:
			inner.WriteString(tok.raw)
		}
	}
	return "", 0
}

// wrapLine splits a long line at word boundaries, preserving leading indentation.
// It is XML-aware: it won't break inside tags (< ... >).
func wrapLine(line string, maxWidth int) []string {
	// Extract leading indentation
	trimmed := strings.TrimLeft(line, " \t")
	indentStr := line[:len(line)-len(trimmed)]
	contIndent := indentStr + "  " // continuation lines get extra indent

	// Split into atoms: tags and words, never broken internally. Whitespace
	// between atoms is collapsed to a single significant space recorded in
	// spaceBefore — including space between a word and an adjacent tag
	// (e.g. "</strong> <code>"), which must survive wrapping.
	atoms := splitAtoms(trimmed)

	var lines []string
	var cur strings.Builder
	cur.WriteString(indentStr)
	curWidth := runeWidth(indentStr)

	for _, a := range atoms {
		atomW := runeWidth(a.text)
		spaceW := 0
		if a.spaceBefore {
			spaceW = 1
		}

		// Wrap before this atom if it would overflow and the current line
		// already holds content beyond its indentation.
		if curWidth+spaceW+atomW > maxWidth && curWidth > runeWidth(contIndent) {
			lines = append(lines, strings.TrimRight(cur.String(), " "))
			cur.Reset()
			cur.WriteString(contIndent)
			curWidth = runeWidth(contIndent)
		} else if a.spaceBefore && curWidth > runeWidth(indentStr) {
			cur.WriteByte(' ')
			curWidth++
		}

		cur.WriteString(a.text)
		curWidth += atomW
	}

	if cur.Len() > 0 {
		final := strings.TrimRight(cur.String(), " ")
		if final != "" {
			lines = append(lines, final)
		}
	}

	if len(lines) == 0 {
		return []string{line}
	}
	return lines
}

// atom is an unbreakable unit of a line — a tag (<...>) or a word — together
// with whether whitespace separated it from the previous atom.
type atom struct {
	text        string
	spaceBefore bool
}

// splitAtoms breaks a line into tags and words, recording the inter-atom
// whitespace as a single space flag. E.g. "Hello <strong>world</strong> end"
// -> {Hello}, { <strong>}, {world}, {</strong>}, { end}.
func splitAtoms(s string) []atom {
	var atoms []atom
	pendingSpace := false
	for len(s) > 0 {
		r := s[0]
		switch {
		case r == ' ' || r == '\t' || r == '\n' || r == '\r':
			pendingSpace = true
			s = s[1:]
		case r == '<':
			end := strings.Index(s, ">")
			if end == -1 {
				atoms = append(atoms, atom{s, pendingSpace})
				return atoms
			}
			// Keep an inline code span whole. A line break inside
			// <code>...</code> would inject the continuation indent into the
			// code's text content, which xml2md then collapses into a spurious
			// space inside the span (e.g. "` app.manifest.lock`").
			if strings.HasPrefix(s, "<code>") || strings.HasPrefix(s, "<code ") {
				if c := strings.Index(s, "</code>"); c != -1 {
					spanEnd := c + len("</code>")
					atoms = append(atoms, atom{s[:spanEnd], pendingSpace})
					pendingSpace = false
					s = s[spanEnd:]
					continue
				}
			}
			atoms = append(atoms, atom{s[:end+1], pendingSpace})
			pendingSpace = false
			s = s[end+1:]
		default:
			// A word runs until the next space or tag start.
			i := strings.IndexAny(s, " \t\n\r<")
			if i == -1 {
				i = len(s)
			}
			atoms = append(atoms, atom{s[:i], pendingSpace})
			pendingSpace = false
			s = s[i:]
		}
	}
	return atoms
}

func runeWidth(s string) int {
	return utf8.RuneCountInString(s)
}

func ensureIndentedLine(buf *strings.Builder, level int, indent string, atLineStart *bool) {
	if !*atLineStart {
		buf.WriteString("\n")
	}
	writeIndentPrefix(buf, level, indent)
	*atLineStart = false
}

func writeIndentPrefix(buf *strings.Builder, level int, indent string) {
	for range level {
		buf.WriteString(indent)
	}
}

// normalizeSelfClose ensures all self-closing tags use the canonical
// "<tag ... />" form (with a single space before the slash).
func normalizeSelfClose(raw string) string {
	if !strings.HasSuffix(raw, "/>") {
		return raw
	}
	inner := strings.TrimSuffix(raw, "/>")
	inner = strings.TrimRight(inner, " \t")
	return inner + " />"
}

// normalizeEntities maps interchangeable entity escapes onto a canonical form.
// Confluence storage produces &quot; and literal apostrophes; goldmark/html
// emits &#34; and &#39;. We normalize all to the Confluence form so verify
// doesn't flag cosmetic differences.
func normalizeEntities(s string) string {
	s = strings.ReplaceAll(s, "&#34;", "&quot;")
	s = strings.ReplaceAll(s, "&#39;", "'")
	s = strings.ReplaceAll(s, "&apos;", "'")
	return s
}

func collapseWS(s string) string {
	var buf strings.Builder
	inWS := false
	for _, r := range s {
		if r == ' ' || r == '\t' || r == '\n' || r == '\r' {
			if !inWS {
				buf.WriteByte(' ')
				inWS = true
			}
		} else {
			buf.WriteRune(r)
			inWS = false
		}
	}
	return buf.String()
}

// Token types for the XML tokenizer.
type tokenKind int

const (
	tokenOpen      tokenKind = iota // <tag ...>
	tokenClose                      // </tag>
	tokenSelfClose                  // <tag .../>
	tokenText                       // plain text
	tokenCDATA                      // <![CDATA[...]]>
	tokenComment                    // <!-- ... -->
)

type token struct {
	kind tokenKind
	raw  string
}

func (t token) tagName() string {
	s := t.raw
	switch t.kind {
	case tokenOpen, tokenSelfClose:
		s = s[1:]
		if strings.HasSuffix(s, "/>") {
			s = s[:len(s)-2]
		} else {
			s = strings.TrimSuffix(s, ">")
		}
		if idx := strings.IndexAny(s, " \t\n"); idx > 0 {
			s = s[:idx]
		}
		return strings.ToLower(s)
	case tokenClose:
		s = s[2:]
		s = strings.TrimSuffix(s, ">")
		return strings.ToLower(strings.TrimSpace(s))
	}
	return ""
}

func tokenize(input string) []token {
	var tokens []token
	i := 0
	for i < len(input) {
		if input[i] == '<' {
			if strings.HasPrefix(input[i:], "<![CDATA[") {
				end := strings.Index(input[i:], "]]>")
				if end == -1 {
					tokens = append(tokens, token{tokenCDATA, input[i:]})
					break
				}
				tokens = append(tokens, token{tokenCDATA, input[i : i+end+3]})
				i += end + 3
				continue
			}
			if strings.HasPrefix(input[i:], "<!--") {
				end := strings.Index(input[i:], "-->")
				if end == -1 {
					tokens = append(tokens, token{tokenComment, input[i:]})
					break
				}
				tokens = append(tokens, token{tokenComment, input[i : i+end+3]})
				i += end + 3
				continue
			}
			end := strings.Index(input[i:], ">")
			if end == -1 {
				tokens = append(tokens, token{tokenText, input[i:]})
				break
			}
			tagStr := input[i : i+end+1]
			if strings.HasPrefix(tagStr, "</") {
				tokens = append(tokens, token{tokenClose, tagStr})
			} else if strings.HasSuffix(tagStr, "/>") {
				tokens = append(tokens, token{tokenSelfClose, tagStr})
			} else {
				tokens = append(tokens, token{tokenOpen, tagStr})
			}
			i += end + 1
		} else {
			end := strings.Index(input[i:], "<")
			if end == -1 {
				tokens = append(tokens, token{tokenText, input[i:]})
				break
			}
			tokens = append(tokens, token{tokenText, input[i : i+end]})
			i += end
		}
	}
	return tokens
}