package main

func alignOutput(src string) string {
	lines := splitLines(src)
	alignFieldTypes(lines)
	alignTrailingComments(lines)
	return joinLines(lines)
}

func splitLines(s string) []string {
	out := []string{:0:64}
	start := int32(0)
	for i := int32(0); i < int32(len(s)); i++ {
		if s[i] == '\n' {
			out = append(out, s[start:i])
			start = i + 1
		}
	}
	if start < int32(len(s)) {
		out = append(out, s[start:])
	}
	return out
}

func joinLines(lines []string) string {
	if len(lines) == 0 { return "" }
	out := []byte{:0:1024}
	for i := int32(0); i < int32(len(lines)); i++ {
		out = append(out, lines[i]...)
		out = append(out, '\n')
	}
	return string(out)
}

func lineIndent(line string) (n int32) {
	for n < int32(len(line)) && line[n] == '\t' { n++ }
	return
}

func trimLeft(s string) string {
	i := int32(0)
	for i < int32(len(s)) && (s[i] == ' ' || s[i] == '\t') { i++ }
	return s[i:]
}

func isKeywordPrefix(s string) (yes bool) {
	kws := [20]string{
		"type ", "func ", "var ", "const ", "for ", "if ", "else ",
		"switch ", "case ", "default:", "select ", "return ",
		"defer ", "go ", "goto ", "break", "continue", "import ",
		"package ", "//",
	}
	for i := int32(0); i < int32(len(kws)); i++ {
		k := kws[i]
		if int32(len(s)) >= int32(len(k)) && s[:len(k)] == k {
			yes = true
			return
		}
	}
	return
}

func isFieldLine(line string) (yes bool) {
	s := trimLeft(line)
	if len(s) == 0 { return }
	if s[0] == '/' { return }
	if !isIdentStart(rune(s[0])) { return }
	if isKeywordPrefix(s) { return }
	spaceIdx := indexByte(s, ' ')
	tabIdx := indexByte(s, '\t')
	sepIdx := spaceIdx
	if sepIdx < 0 || (tabIdx >= 0 && tabIdx < sepIdx) { sepIdx = tabIdx }
	if sepIdx < 0 { return }
	rest := trimLeft(s[sepIdx:])
	if len(rest) == 0 { return }
	if rest[0] == '=' || rest[0] == '+' || rest[0] == '-' || rest[0] == ':' ||
		rest[0] == '|' || rest[0] == '(' { return }
	yes = true
	return
}

func isIdentStart(ch rune) (yes bool) {
	if 'a' <= ch && ch <= 'z' || 'A' <= ch && ch <= 'Z' || ch == '_' {
		yes = true
	}
	return
}

func indexByte(s string, b uint8) (idx int32) {
	for idx = 0; idx < int32(len(s)); idx++ {
		if s[idx] == b { return }
	}
	idx = -1
	return
}

func fieldNameEnd(line string) (pos int32) {
	s := trimLeft(line)
	indent := int32(len(line)) - int32(len(s))
	pos = indent
	for {
		for pos < int32(len(line)) && line[pos] != ' ' && line[pos] != '\t' && line[pos] != ',' { pos++ }
		if pos < int32(len(line)) && line[pos] == ',' {
			pos++
			for pos < int32(len(line)) && line[pos] == ' ' { pos++ }
			continue
		}
		break
	}
	return
}

func commentStart(line string) (pos int32) {
	pos = -1
	inStr := false
	esc := false
	for i := int32(0); i < int32(len(line)) - 1; i++ {
		ch := line[i]
		if esc { esc = false; continue }
		if ch == '\\' { esc = true; continue }
		if ch == '"' || ch == '\'' || ch == '`' { inStr = !inStr; continue }
		if inStr { continue }
		if ch == '/' && line[i+1] == '/' {
			pos = i
			return
		}
	}
	return
}

// --- field type alignment ---

func alignFieldTypes(lines []string) {
	i := int32(0)
	n := int32(len(lines))
	for i < n {
		gStart, gEnd := findFieldGroup(lines, i, n)
		if gEnd > gStart {
			alignGroup(lines, gStart, gEnd)
			i = gEnd
		} else {
			i++
		}
	}
}

func findFieldGroup(lines []string, from int32, n int32) (gStart int32, gEnd int32) {
	gStart = from
	for gStart < n && !isFieldLine(lines[gStart]) { gStart++ }
	if gStart >= n { gEnd = gStart; return }
	indent := lineIndent(lines[gStart])
	gEnd = gStart + 1
	blanks := int32(0)
	for gEnd < n {
		line := lines[gEnd]
		if len(line) == 0 || len(trimLeft(line)) == 0 {
			blanks++
			if blanks >= 2 { break }
			gEnd++
			continue
		}
		blanks = 0
		if lineIndent(line) != indent { break }
		if !isFieldLine(line) { break }
		gEnd++
	}
	if gEnd - gStart < 2 { gEnd = gStart }
	return
}

func alignGroup(lines []string, from int32, to int32) {
	maxNameW := int32(0)
	for i := from; i < to; i++ {
		line := lines[i]
		if len(trimLeft(line)) == 0 { continue }
		ne := fieldNameEnd(line)
		indent := lineIndent(line)
		nameW := ne - indent
		if nameW > maxNameW { maxNameW = nameW }
	}
	for i := from; i < to; i++ {
		line := lines[i]
		if len(trimLeft(line)) == 0 { continue }
		ne := fieldNameEnd(line)
		indent := lineIndent(line)
		nameW := ne - indent
		pad := maxNameW - nameW
		if pad <= 0 { continue }
		rest := line[ne:]
		rest = trimLeft(rest)
		out := []byte{:0:int32(len(line)) + pad}
		out = append(out, line[:ne]...)
		for j := int32(0); j <= pad; j++ { out = append(out, ' ') }
		out = append(out, rest...)
		lines[i] = string(out)
	}
}

// --- trailing comment alignment ---

func alignTrailingComments(lines []string) {
	i := int32(0)
	n := int32(len(lines))
	for i < n {
		gStart := i
		for gStart < n {
			if commentStart(lines[gStart]) >= 0 { break }
			gStart++
		}
		if gStart >= n { break }
		gEnd := gStart + 1
		blanks := int32(0)
		for gEnd < n {
			line := lines[gEnd]
			if len(line) == 0 || len(trimLeft(line)) == 0 {
				blanks++
				if blanks >= 2 { break }
				gEnd++
				continue
			}
			blanks = 0
			if commentStart(line) < 0 { break }
			gEnd++
		}
		if gEnd - gStart >= 2 {
			alignCommentGroup(lines, gStart, gEnd)
		}
		i = gEnd
	}
}

func alignCommentGroup(lines []string, from int32, to int32) {
	maxCode := int32(0)
	for i := from; i < to; i++ {
		cs := commentStart(lines[i])
		if cs < 0 { continue }
		codeW := cs
		for codeW > 0 && (lines[i][codeW-1] == ' ' || lines[i][codeW-1] == '\t') {
			codeW--
		}
		if codeW > maxCode { maxCode = codeW }
	}
	for i := from; i < to; i++ {
		cs := commentStart(lines[i])
		if cs < 0 { continue }
		codeW := cs
		for codeW > 0 && (lines[i][codeW-1] == ' ' || lines[i][codeW-1] == '\t') {
			codeW--
		}
		cmt := lines[i][cs:]
		out := []byte{:0:maxCode + 1 + int32(len(cmt))}
		out = append(out, lines[i][:codeW]...)
		pad := maxCode - codeW + 1
		for j := int32(0); j < pad; j++ { out = append(out, ' ') }
		out = append(out, cmt...)
		lines[i] = string(out)
	}
}