package pglex

import (
	"bytes"
	"strconv"
	"strings"

	"code.justin.tv/release/trace/pgparse"
)

func Normalize(sql string) string {
	tkn := pgparse.NewStringTokenizer(sql)
	w := bytes.NewBuffer(make([]byte, 0, 1024))

	var prevTk int
	var prevBuf []byte

	for {
		tk, buf := tkn.Scan()

		if tk == 0 && buf == nil {
			// EOF
			break
		}

		if buf == nil {
			// symbol
			switch tk {
			case '.', ',', ')', '[':
				// some symbols need no leading space
			default:
				if prevBuf != nil {
					// insert space between non-symbols and symbols
					w.WriteByte(' ')
				}
			}
			w.WriteRune(rune(tk))
		} else if operator := pgparse.OperatorFromToken(tk); operator != "" {
			//operator which needs special consideration
			if tk != pgparse.TYPECAST {
				w.WriteByte(' ')
			}

			w.WriteString(operator)
		} else if keyword := pgparse.KeywordFromToken(tk); keyword != "" {
			// SQL keyword
			stringPrefixSpace(prevTk, w)
			w.WriteString(strings.ToUpper(keyword))
		} else {
			// non-symbol, non-keyword
			switch tk {
			case pgparse.STRING, pgparse.NUMBER, pgparse.VALUE_ARG, pgparse.LIST_ARG:
				// data must be removed from the query
				stringPrefixSpace(prevTk, w)
				w.WriteByte('?')
				// consider this a symbol when calculating if a space should follow
				buf = nil
			case pgparse.ID:
				// e.g. table or column name - leave these in (quoted)
				stringPrefixSpace(prevTk, w)
				w.WriteString(strconv.Quote(string(buf)))
			case pgparse.COMMENT:
			case pgparse.LEX_ERROR:
				stringPrefixSpace(prevTk, w)
				w.WriteString("'lex error'")
				// we can't assume we're correctly normalizing the query
				return w.String()
			default:
				// allow debugging of unhandled token types
				stringPrefixSpace(prevTk, w)
				w.WriteString("'token-")
				w.WriteString(strconv.FormatInt(int64(tk), 10))
				w.WriteString("'")
			}
		}

		prevTk, prevBuf = tk, buf
	}

	return w.String()
}

func stringPrefixSpace(prevTk int, w *bytes.Buffer) {
	switch prevTk {
	case 0, '(', '.', '[', pgparse.TYPECAST:
		// no space required at the beginning of the string, or after particular symbols
	default:
		// insert space before keywords
		w.WriteByte(' ')
	}
}
