package main

import (
	"os"
)

const maxLine = int32(80)

type Printer struct {
	toks      []FTok
	out       []byte
	pos       int32
	indent    int32
	col       int32
	atBOL     bool
	errs      []string
	caseDepth int32
}

func newPrinter(toks []FTok) (p *Printer) {
	p = &Printer{
		toks:   toks,
		out:    []byte{:0:4096},
		errs:   []string{:0:4},
		pos:    0,
		indent: 0,
		col:    0,
		atBOL:  true,
	}
	return
}

func (p *Printer) errorf(line uint32, msg string) {
	s := int32(line).String() | ":" | msg
	p.errs = append(p.errs, s)
}

func (p *Printer) emit(s string) {
	p.out = append(p.out, s...)
	for i := int32(0); i < int32(len(s)); i++ {
		if s[i] == '\n' {
			p.col = 0
			p.atBOL = true
		} else if s[i] == '\t' {
			p.col += 4 - (p.col % 4)
		} else {
			p.col++
		}
	}
}

func (p *Printer) newline() { p.emit("\n"); p.atBOL = true }

func (p *Printer) writeIndent() {
	for i := int32(0); i < p.indent; i++ { p.emit("\t") }
	p.atBOL = false
}

func (p *Printer) space() { p.emit(" ") }

func (p *Printer) cur() (t *FTok) { t = &p.toks[p.pos]; return }

func (p *Printer) peek(off int32) (t *FTok) {
	idx := p.pos + off
	if idx < 0 { idx = 0 }
	if idx >= int32(len(p.toks)) { idx = int32(len(p.toks)) - 1 }
	t = &p.toks[idx]
	return
}

func (p *Printer) advance() {
	if p.pos < int32(len(p.toks))-1 { p.pos++ }
}

// --- matching ---

func matchBrace(toks []FTok, from int32) (end int32) {
	depth := int32(1)
	i := from + 1
	for i < int32(len(toks)) {
		tt := toks[i].tok
		if tt == TokLbrace { depth++ } else if tt == TokRbrace {
			depth--
			if depth == 0 { end = i; return }
		} else if tt == TokEOF { end = i; return }
		i++
	}
	end = int32(len(toks)) - 1
	return
}

func matchBrack(toks []FTok, from int32) (end int32) {
	depth := int32(1)
	i := from + 1
	for i < int32(len(toks)) {
		tt := toks[i].tok
		if tt == TokLbrack { depth++ } else if tt == TokRbrack {
			depth--
			if depth == 0 { end = i; return }
		} else if tt == TokEOF { end = i; return }
		i++
	}
	end = int32(len(toks)) - 1
	return
}

func matchParen(toks []FTok, from int32) (end int32) {
	depth := int32(1)
	i := from + 1
	for i < int32(len(toks)) {
		tt := toks[i].tok
		if tt == TokLparen { depth++ } else if tt == TokRparen {
			depth--
			if depth == 0 { end = i; return }
		} else if tt == TokEOF { end = i; return }
		i++
	}
	end = int32(len(toks)) - 1
	return
}

// --- width measurement ---

func singleLineWidth(toks []FTok, from int32, to int32) (w int32) {
	w = 0
	for i := from; i <= to; i++ {
		t := &toks[i]
		if t.tok == TokComment { w = maxLine + 1; return }
		if t.tok == TokSemi { w += 2; continue }
		txt := tokText(t)
		w += int32(len(txt))
		if i < to {
			next := &toks[i+1]
			if needSpace(t, next) { w++ }
		}
	}
	return
}

func parenFitsLine(toks []FTok, openIdx int32, startCol int32) (yes bool) {
	end := matchParen(toks, openIdx)
	w := int32(1)
	for i := openIdx + 1; i <= end; i++ {
		t := &toks[i]
		if t.tok == TokSemi { w += 2; continue }
		if t.tok == TokComment { return }
		w += int32(len(tokText(t)))
		if i < end {
			next := &toks[i+1]
			if needSpace(t, next) { w++ }
		}
	}
	yes = startCol + w <= maxLine
	return
}

func hasInnerBrace(toks []FTok, from int32, to int32) (yes bool) {
	for i := from; i <= to; i++ {
		if toks[i].tok == TokLbrace { yes = true; return }
	}
	return
}

// --- binary/unary disambiguation ---

func isBinaryMulStar(toks []FTok, idx int32) (yes bool) {
	if idx <= 0 { return }
	prev := toks[idx-1].tok
	if prev != TokName && prev != TokLit && prev != TokRparen &&
		prev != TokRbrack && prev != TokRbrace { return }
	if idx < 2 { yes = true; return }
	pp := toks[idx-2].tok
	if pp == TokVar || pp == TokConst || pp == TokType || pp == TokComma ||
		pp == TokFunc || pp == TokChan || pp == TokMap || pp == TokReturn ||
		pp == TokAssign || pp == TokDefine || pp == TokArrow ||
		pp == TokColon || pp == TokLparen || pp == TokLbrack ||
		pp == TokCase || pp == TokSemi || pp == TokLbrace {
		return
	}
	yes = true
	return
}

func isBinaryArith(toks []FTok, idx int32) (yes bool) {
	if idx <= 0 { return }
	prev := toks[idx-1].tok
	if prev == TokName || prev == TokLit || prev == TokRparen ||
		prev == TokRbrack || prev == TokRbrace {
		yes = true
	}
	return
}

// --- brace spacing ---

func braceNeedsSpace(toks []FTok, braceIdx int32) (yes bool) {
	if braceIdx < 1 { return }
	prev := toks[braceIdx-1].tok
	if prev != TokName && prev != TokRparen { return }
	if prev == TokRparen { yes = true; return }
	for i := braceIdx - 1; i >= 0; i-- {
		tt := toks[i].tok
		if tt == TokFunc || tt == TokIf || tt == TokFor || tt == TokElse ||
			tt == TokSelect || tt == TokSwitch {
			yes = true; return
		}
		if tt == TokSemi || tt == TokLbrace || tt == TokRbrace || tt == TokEOF { return }
	}
	return
}

func isSliceColonCtx(toks []FTok, idx int32) (yes bool) {
	bdepth := int32(0)
	pdepth := int32(0)
	for i := idx - 1; i >= 0; i-- {
		tt := toks[i].tok
		if tt == TokRbrack { bdepth++ } else if tt == TokLbrack {
			if bdepth == 0 { yes = true; return }
			bdepth--
		} else if tt == TokRbrace { pdepth++ } else if tt == TokLbrace {
			if pdepth == 0 { yes = true; return }
			pdepth--
		} else if tt == TokSemi { return }
	}
	return
}

// --- spacing ---

func needSpaceCtx(toks []FTok, curIdx int32, prev *FTok, cur *FTok) (yes bool) {
	if cur.tok == TokLbrack && prev.tok == TokName && curIdx+1 < int32(len(toks)) {
		next := toks[curIdx+1].tok
		if next == TokRbrack {
			rbIdx := curIdx + 2
			if rbIdx < int32(len(toks)) {
				after := toks[rbIdx].tok
				if after == TokName || after == TokStar || after == TokLbrack ||
					after == TokFunc || after == TokMap || after == TokChan ||
					after == TokStruct || after == TokInterface {
					yes = true; return
				}
			}
		}
		if next == TokLit {
			end := matchBrack(toks, curIdx)
			if end+1 < int32(len(toks)) {
				after := toks[end+1].tok
				if after == TokName || after == TokStar || after == TokLbrack ||
					after == TokFunc || after == TokMap || after == TokChan ||
					after == TokStruct || after == TokInterface {
					yes = true; return
				}
			}
		}
	}
	if cur.tok == TokLbrace && braceNeedsSpace(toks, curIdx) {
		yes = true; return
	}
	if prev.tok == TokStar && isBinaryMulStar(toks, curIdx-1) {
		yes = true; return
	}
	if cur.tok == TokStar && isBinaryMulStar(toks, curIdx) {
		yes = true; return
	}
	if prev.tok == TokOp && (prev.lit == "-" || prev.lit == "+" || prev.lit == "^") {
		if !isBinaryArith(toks, curIdx-1) { return }
	}
	if prev.tok == TokOp && (prev.lit == "&" || prev.lit == "!") {
		if !isBinaryArith(toks, curIdx-1) { return }
	}
	if cur.tok == TokColon && isSliceColonCtx(toks, curIdx) { return }
	if prev.tok == TokColon && isSliceColonCtx(toks, curIdx-1) { return }
	yes = needSpaceBefore(prev, cur)
	return
}

func needSpaceBefore(prev *FTok, cur *FTok) (yes bool) {
	if cur.tok == TokLbrace {
		pt := prev.tok
		if pt == TokStruct || pt == TokInterface || pt == TokRbrack ||
			pt == TokName || pt == TokStar { return }
		yes = true; return
	}
	yes = needSpace(prev, cur)
	return
}

func needSpace(prev *FTok, next *FTok) (yes bool) {
	pt := prev.tok
	nt := next.tok

	if pt == TokComment || nt == TokComment { yes = true; return }
	if nt == TokIncOp { return }

	if nt == TokLparen {
		if pt == TokName || pt == TokRparen || pt == TokRbrack { return }
		yes = true; return
	}
	if nt == TokLbrack {
		if pt == TokName || pt == TokRparen || pt == TokRbrack || pt == TokLit { return }
	}
	if pt == TokRbrack {
		if nt == TokName || nt == TokStar || nt == TokLbrack || nt == TokMap ||
			nt == TokStruct || nt == TokInterface || nt == TokFunc || nt == TokChan { return }
	}
	if pt == TokLparen || pt == TokLbrack || nt == TokRparen || nt == TokRbrack { return }
	if nt == TokComma || nt == TokSemi || nt == TokColon { return }
	if pt == TokComma || pt == TokSemi { yes = true; return }
	if pt == TokColon { yes = true; return }
	if nt == TokDot || pt == TokDot { return }

	if pt == TokStar {
		if nt == TokName || nt == TokLparen || nt == TokStar || nt == TokLbrack { return }
		yes = true; return
	}
	if nt == TokStar {
		if pt == TokLparen || pt == TokComma || pt == TokLbrack ||
			pt == TokFunc || pt == TokChan || pt == TokMap ||
			pt == TokArrow || pt == TokAssign || pt == TokDefine ||
			pt == TokReturn || pt == TokCase || pt == TokOp ||
			pt == TokColon || pt == TokDotDotDot { return }
		yes = true; return
	}

	if pt == TokOp && prev.lit == "!" { return }

	if nt == TokLbrace {
		pt2 := prev.tok
		if pt2 == TokStruct || pt2 == TokInterface || pt2 == TokRbrack || pt2 == TokName { return }
		yes = true; return
	}
	if pt == TokLbrace || nt == TokRbrace { yes = true; return }
	if pt == TokOp || pt == TokAssign || pt == TokDefine || pt == TokAssignOp ||
		pt == TokArrow { yes = true; return }
	if nt == TokOp || nt == TokAssign || nt == TokDefine || nt == TokAssignOp ||
		nt == TokArrow { yes = true; return }
	if pt == TokIncOp { yes = true; return }
	if pt == TokDotDotDot || nt == TokDotDotDot { return }

	yes = true
	return
}

// --- predicates ---

func isEmptyBrace(toks []FTok, bracePos int32) (yes bool) {
	end := matchBrace(toks, bracePos)
	for i := bracePos + 1; i < end; i++ {
		if toks[i].tok != TokSemi { return }
	}
	yes = true; return
}

func fitsOneLine(toks []FTok, bracePos int32, preambleCol int32) (yes bool) {
	end := matchBrace(toks, bracePos)
	innerFrom := bracePos + 1
	innerTo := end - 1
	if innerTo < innerFrom { yes = true; return }
	for i := innerFrom; i <= innerTo; i++ {
		if toks[i].tok == TokComment { return }
	}
	if hasInnerBrace(toks, innerFrom, innerTo) { return }
	w := singleLineWidth(toks, innerFrom, innerTo)
	total := preambleCol + 2 + w + 2
	yes = total <= maxLine
	return
}

func isGroupKeyword(tok Token) (yes bool) {
	if tok == TokImport || tok == TokConst || tok == TokVar || tok == TokType { yes = true }
	return
}

func isImportContext(p *Printer) (yes bool) {
	for i := p.pos - 1; i >= 0; i-- {
		t := &p.toks[i]
		if t.tok == TokImport { yes = true; return }
		if t.tok == TokSemi || t.tok == TokRbrace || t.tok == TokLbrace { return }
	}
	return
}

func isLabel(toks []FTok, colonIdx int32) (yes bool) {
	if colonIdx < 1 { return }
	if toks[colonIdx-1].tok != TokName { return }
	if colonIdx < 2 { yes = true; return }
	pp := toks[colonIdx-2].tok
	if pp == TokSemi || pp == TokLbrace || pp == TokColon { yes = true }
	return
}

func isSliceLitBrace(toks []FTok, bracePos int32) (yes bool) {
	end := matchBrace(toks, bracePos)
	for i := bracePos + 1; i < end; i++ {
		if toks[i].tok == TokSemi { continue }
		if toks[i].tok == TokColon { yes = true }
		return
	}
	return
}

// --- bare return type diagnostic ---

func checkBareReturns(toks []FTok, p *Printer) {
	i := int32(0)
	for i < int32(len(toks)) {
		if toks[i].tok == TokFunc {
			checkFuncSig(toks, i, p)
		}
		i++
	}
}

func checkFuncSig(toks []FTok, funcIdx int32, p *Printer) {
	i := funcIdx + 1
	n := int32(len(toks))
	if i >= n { return }
	if toks[i].tok == TokLparen {
		i = matchParen(toks, i) + 1
	}
	if i >= n { return }
	if toks[i].tok != TokName { return }
	i++
	if i >= n { return }
	if toks[i].tok != TokLparen { return }
	i = matchParen(toks, i) + 1
	if i >= n { return }
	tt := toks[i].tok
	if tt == TokLbrace || tt == TokSemi { return }
	if tt == TokLparen { return }
	p.errorf(toks[funcIdx].line, " bare return type (wrap in parens with name)")
}

// --- line-break after comma at 80+ ---

func (p *Printer) emitSpaceOrBreak(prev *FTok) {
	if prev.tok == TokComma && p.col >= maxLine {
		p.newline()
		p.writeIndent()
	} else {
		p.space()
	}
}

// --- paren group wrapping ---

func (p *Printer) formatParenGroup() {
	end := matchParen(p.toks, p.pos)
	if parenFitsLine(p.toks, p.pos, p.col) {
		p.emit("(")
		p.advance()
		for p.pos < end {
			t := p.cur()
			if t.tok == TokSemi { p.advance(); continue }
			prev := p.peek(-1)
			if needSpaceCtx(p.toks, p.pos, prev, t) { p.space() }
			p.emit(tokText(t))
			p.advance()
		}
		p.emit(")")
		p.advance()
		return
	}
	p.emit("(")
	p.advance()
	p.indent++

	needNL2 := true
	for p.pos < end {
		t := p.cur()
		if t.tok == TokSemi { p.advance(); continue }
		if t.tok == TokComment {
			p.formatComment(t)
			p.advance()
			continue
		}
		if t.tok == TokComma {
			p.emit(",")
			p.advance()
			needNL2 = true
			continue
		}
		if needNL2 {
			p.newline()
			p.writeIndent()
			needNL2 = false
		} else {
			prev := p.peek(-1)
			if needSpaceCtx(p.toks, p.pos, prev, t) { p.space() }
		}
		if t.tok == TokLparen {
			p.formatParenGroup()
		} else {
			p.emit(tokText(t))
			p.advance()
		}
	}

	p.indent--
	p.newline()
	p.writeIndent()
	p.emit(")")
	p.advance()
}

// --- main format ---

func (p *Printer) format() {
	for p.pos < int32(len(p.toks)) {
		t := p.cur()
		if t.tok == TokEOF { break }
		if t.tok == TokComment { p.formatComment(t); p.advance(); continue }
		if t.tok == TokSemi { p.advance(); continue }

		if t.tok == TokElse && p.pos > 0 {
			prev := p.peek(-1)
			if prev.tok == TokRbrace {
				p.space()
				p.emit(tokText(t))
				p.advance()
				continue
			}
		}

		if t.nlBefore && p.pos > 0 {
			p.ensureNewline()
			if t.blankBef { p.newline() }
			p.writeIndent()
		} else if p.atBOL {
			p.writeIndent()
		} else if p.pos > 0 {
			prev := p.peek(-1)
			if needSpaceCtx(p.toks, p.pos, prev, t) { p.space() }
		}

		if t.tok == TokLbrace {
			p.formatBrace()
		} else if t.tok == TokLparen && p.pos > 0 && isGroupKeyword(p.peek(-1).tok) {
			p.formatGroupBlock()
		} else if t.tok == TokLparen && !parenFitsLine(p.toks, p.pos, p.col) {
			p.formatParenGroup()
		} else {
			p.emit(tokText(t))
			p.advance()
		}
	}

	if len(p.out) > 0 && p.out[len(p.out)-1] != '\n' { p.newline() }
}

func (p *Printer) ensureNewline() {
	if len(p.out) > 0 && p.out[len(p.out)-1] != '\n' { p.newline() }
}

func (p *Printer) formatComment(t *FTok) {
	if t.nlBefore && p.pos > 0 {
		p.ensureNewline()
		if t.blankBef { p.newline() }
		p.writeIndent()
	} else if p.atBOL {
		p.writeIndent()
	} else {
		p.space()
	}
	p.emit(t.lit)
}

// --- brace formatting ---

func (p *Printer) formatBrace() {
	if fitsOneLine(p.toks, p.pos, p.col) && !isImportContext(p) {
		p.formatBraceSingleLine()
	} else {
		p.formatBraceMultiLine()
	}
}

func (p *Printer) formatBraceSingleLine() {
	end := matchBrace(p.toks, p.pos)
	p.emit("{")
	p.advance()

	if isEmptyBrace(p.toks, p.pos-1) {
		for p.pos < end { p.advance() }
		p.emit("}")
		p.advance()
		return
	}

	sliceLit := isSliceLitBrace(p.toks, p.pos-1)
	if !sliceLit { p.space() }
	first := true
	for p.pos < end {
		t := p.cur()
		if t.tok == TokSemi {
			next := p.peek(1)
			if next.tok != TokRbrace && p.pos+1 < end { p.emit("; ") }
			p.advance()
			first = true
			continue
		}
		if !first {
			prev := p.peek(-1)
			if prev.tok != TokSemi && needSpaceCtx(p.toks, p.pos, prev, t) { p.space() }
		}
		first = false
		p.emit(tokText(t))
		p.advance()
	}
	if !sliceLit { p.space() }
	p.emit("}")
	p.advance()
}

func (p *Printer) formatBraceMultiLine() {
	end := matchBrace(p.toks, p.pos)
	p.emit("{")
	p.advance()
	p.indent++

	caseIndented := false
	caseAwait := false
	needNL := true
	for p.pos < end {
		t := p.cur()

		if t.tok == TokSemi {
			if t.lit == ";" {
				p.emit(";")
				needNL = false
			} else {
				needNL = true
			}
			p.advance()
			continue
		}

		if t.tok == TokComment {
			p.formatComment(t)
			needNL = true
			p.advance()
			continue
		}

		if t.tok == TokElse {
			p.space()
			p.emit(tokText(t))
			p.advance()
			needNL = false
			continue
		}

		if (t.tok == TokCase || t.tok == TokDefault) && needNL {
			if caseIndented { p.indent--; caseIndented = false }
			caseAwait = true
			p.ensureNewline()
			if t.blankBef { p.newline() }
			p.writeIndent()
			needNL = false
			p.emit(tokText(t))
			p.advance()
			continue
		}

		if needNL {
			p.ensureNewline()
			if t.blankBef { p.newline() }
			p.writeIndent()
			needNL = false
		} else {
			prev := p.peek(-1)
			if prev.tok == TokComma && p.col >= maxLine {
				p.newline()
				p.writeIndent()
			} else if needSpaceCtx(p.toks, p.pos, prev, t) {
				p.space()
			}
		}

		if t.tok == TokColon && caseAwait {
			p.emit(":")
			p.advance()
			caseAwait = false
			caseIndented = true
			p.indent++
			needNL = true
		} else if t.tok == TokColon && isLabel(p.toks, p.pos) {
			p.emit(":")
			p.advance()
			needNL = true
		} else if t.tok == TokLbrace {
			p.formatBrace()
			needNL = true
		} else if t.tok == TokLparen && p.pos > 0 && isGroupKeyword(p.peek(-1).tok) {
			p.formatGroupBlock()
		} else if t.tok == TokLparen && !parenFitsLine(p.toks, p.pos, p.col) {
			p.formatParenGroup()
		} else {
			p.emit(tokText(t))
			p.advance()
		}
	}

	if caseIndented { p.indent-- }
	p.indent--
	p.ensureNewline()
	p.writeIndent()
	p.emit("}")
	p.advance()
}

// --- group block (import/const/var/type) ---

func (p *Printer) formatGroupBlock() {
	end := matchParen(p.toks, p.pos)
	p.emit("(")
	p.advance()
	p.indent++

	needNL := true
	for p.pos < end {
		t := p.cur()
		if t.tok == TokSemi { needNL = true; p.advance(); continue }
		if t.tok == TokComment {
			p.formatComment(t)
			needNL = true
			p.advance()
			continue
		}
		if needNL {
			p.ensureNewline()
			if t.blankBef { p.newline() }
			p.writeIndent()
			needNL = false
		} else {
			prev := p.peek(-1)
			if needSpaceCtx(p.toks, p.pos, prev, t) { p.space() }
		}
		p.emit(tokText(t))
		p.advance()
	}

	p.indent--
	p.ensureNewline()
	p.writeIndent()
	p.emit(")")
	p.advance()
}

// --- entry ---

func formatSource(src string) (out string, errs []string) {
	toks := tokenize(src)
	p := newPrinter(toks)
	checkBareReturns(toks, p)
	p.format()
	out = alignOutput(string(p.out))
	errs = p.errs
	return
}

func reportErrors(path string, errs []string) {
	for i := int32(0); i < int32(len(errs)); i++ {
		os.Stderr().Write([]byte(path | ":" | errs[i] | "\n"))
	}
}