package main

import (
	"fmt"
	"github.com/DataDog/go-sqllexer"
)

type TokenGroupType int

const (
	TOKEN_OTHER   TokenGroupType = iota
	TOKEN_ACT                    // SELECT, CREATE, etc
	TOKEN_FILTER                 // WHERE, AND, etc
	TOKEN_RELATE                 // joins
	TOKEN_DEFINE                 // column data types, etc
	TOKEN_CONTROL                // sorting, etc
)

var tokenGroupMap = map[sqllexer.TokenType]TokenGroupType{
	sqllexer.ERROR:                  TOKEN_OTHER,
	sqllexer.EOF:                    TOKEN_OTHER,
	sqllexer.SPACE:                  TOKEN_OTHER,
	sqllexer.STRING:                 TOKEN_DEFINE,
	sqllexer.INCOMPLETE_STRING:      TOKEN_DEFINE,
	sqllexer.NUMBER:                 TOKEN_DEFINE,
	sqllexer.IDENT:                  TOKEN_ACT,    // for table/column names
	sqllexer.QUOTED_IDENT:           TOKEN_RELATE, // for joins
	sqllexer.OPERATOR:               TOKEN_FILTER,
	sqllexer.WILDCARD:               TOKEN_ACT,
	sqllexer.COMMENT:                TOKEN_OTHER,
	sqllexer.MULTILINE_COMMENT:      TOKEN_OTHER,
	sqllexer.PUNCTUATION:            TOKEN_OTHER,
	sqllexer.DOLLAR_QUOTED_FUNCTION: TOKEN_DEFINE,
	sqllexer.DOLLAR_QUOTED_STRING:   TOKEN_DEFINE,
	sqllexer.POSITIONAL_PARAMETER:   TOKEN_FILTER,
	sqllexer.BIND_PARAMETER:         TOKEN_FILTER,
	sqllexer.FUNCTION:               TOKEN_DEFINE,
	sqllexer.SYSTEM_VARIABLE:        TOKEN_DEFINE,
	sqllexer.UNKNOWN:                TOKEN_OTHER,
	sqllexer.COMMAND:                TOKEN_ACT,
	sqllexer.KEYWORD:                TOKEN_ACT,
	sqllexer.JSON_OP:                TOKEN_DEFINE,
	sqllexer.BOOLEAN:                TOKEN_DEFINE,
	sqllexer.NULL:                   TOKEN_DEFINE,
	sqllexer.PROC_INDICATOR:         TOKEN_OTHER,
	sqllexer.CTE_INDICATOR:          TOKEN_OTHER,
	sqllexer.ALIAS_INDICATOR:        TOKEN_OTHER,
}

func GetTokenGroupType(tokenType sqllexer.TokenType) (TokenGroupType, bool, bool) {
	group, exists := tokenGroupMap[tokenType]
	isStart := false
	isEnd := false

	if !exists {
		group = TOKEN_OTHER
	}

	if group == TOKEN_ACT {
		isStart = true
	}

	if tokenType == sqllexer.EOF ||
		(group == TOKEN_OTHER && tokenType == sqllexer.PUNCTUATION) {
		isEnd = true
	}

	return group, isStart, isEnd
}

func main() {
	query := "SELECT * FROM users \n WHERE id = something AND SELECT;"
	lexer := sqllexer.New(query)
	for {
		token := lexer.Scan()
		tokenGroup, isStart, isEnd := GetTokenGroupType(token.Type)
		fmt.Println(token.Value, token.Type, tokenGroup, isStart, isEnd)

		if token.Type == sqllexer.EOF {
			break
		}
	}
}