package soql

import (
	"strings"
	"text/scanner"
)

// TokenType represents the type of a token
type TokenType int

const (
	TokenError TokenType = iota
	TokenEOF
	TokenIdentifier
	TokenKeyword
	TokenString
	TokenNumber
	TokenOperator
	TokenPunctuation
)

// Token represents a lexical token
type Token struct {
	Type    TokenType
	Literal string
	Pos     int
}

// Lexer tokenizes SOQL queries
type Lexer struct {
	s   scanner.Scanner
	buf string
}

func NewLexer(input string) *Lexer {
	l := &Lexer{buf: input}
	l.s.Init(strings.NewReader(input))
	// Disable default string/char scanning to handle SOQL quotes manually
	l.s.Mode = scanner.ScanIdents | scanner.ScanFloats | scanner.ScanInts
	return l
}

func (l *Lexer) NextToken() Token {
	// 1. Skip Whitespace manually to safely Peek
	l.skipWhitespace()

	pos := l.s.Position.Offset
	ch := l.s.Peek()

	if ch == scanner.EOF {
		return Token{Type: TokenEOF, Literal: "", Pos: pos}
	}

	// 2. Handle Quotes (Strings)
	if ch == '\'' {
		return l.scanString('\'')
	}
	if ch == '"' {
		return l.scanString('"')
	}

	// 3. Handle Multi-char Operators (!=, <=, >=)
	if ch == '!' {
		if l.isNext('=') {
			return Token{Type: TokenOperator, Literal: "!=", Pos: pos}
		}
	}
	if ch == '<' || ch == '>' {
		op := string(ch)
		l.s.Next() // consume first char
		if l.s.Peek() == '=' {
			l.s.Next() // consume '='
			return Token{Type: TokenOperator, Literal: op + "=", Pos: pos}
		}
		return Token{Type: TokenOperator, Literal: op, Pos: pos}
	}

	// 4. Scan Token
	tok := l.s.Scan()
	lit := l.s.TokenText()

	// Handle scanner.EOF case if Scan returns it
	if tok == scanner.EOF {
		return Token{Type: TokenEOF, Literal: "", Pos: pos}
	}

	if tok == scanner.Float || tok == scanner.Int {
		return Token{Type: TokenNumber, Literal: lit, Pos: pos}
	}

	if tok == scanner.Ident {
		if isKeyword(lit) {
			return Token{Type: TokenKeyword, Literal: strings.ToUpper(lit), Pos: pos}
		}
		if isOperator(lit) {
			return Token{Type: TokenOperator, Literal: strings.ToUpper(lit), Pos: pos}
		}
		// Handle "test@example.com" sticking together if scanner parses it?
		// Default scanner separates at '@'.
		// We rely on Parser to recombine or Lexer to be smarter?
		// For now, let's stick to standard behavior.
		return Token{Type: TokenIdentifier, Literal: lit, Pos: pos}
	}

	// Punctuation or Operators
	if isOperator(lit) {
		return Token{Type: TokenOperator, Literal: lit, Pos: pos}
	}

	return Token{Type: TokenPunctuation, Literal: lit, Pos: pos}
}

func (l *Lexer) skipWhitespace() {
	for {
		ch := l.s.Peek()
		if ch == ' ' || ch == '\t' || ch == '\n' || ch == '\r' {
			l.s.Next() // consume
			continue
		}
		break
	}
}

func (l *Lexer) scanString(quote rune) Token {
	pos := l.s.Position.Offset
	l.s.Next() // consume opening quote
	var sb strings.Builder
	for {
		ch := l.s.Next()
		if ch == quote {
			// Check for double-quote escaping (SQL standard: 'It''s')
			if l.s.Peek() == quote {
				l.s.Next() // consume 2nd quote
				sb.WriteRune(quote)
				continue
			}
			break
		}
		if ch == scanner.EOF {
			break
		}
		if ch == '\\' {
			// Handle backslash escape (SOQL supports this too?)
			// Yes, 'O\'Reilly'
			next := l.s.Next()
			if next != scanner.EOF {
				sb.WriteRune(next)
			}
			continue
		}
		sb.WriteRune(ch)
	}
	return Token{Type: TokenString, Literal: sb.String(), Pos: pos}
}

func (l *Lexer) isNext(expected rune) bool {
	l.s.Next() // consume current
	if l.s.Peek() == expected {
		l.s.Next() // consume expected
		return true
	}
	// Warning: We consumed 'current' (e.g. '!') but not expected.
	// This is slightly destructive if not expected?
	// But for '!' in SOQL, '!x' is invalid anyway unless '!='.
	return false
}

func isKeyword(lit string) bool {
	upper := strings.ToUpper(lit)
	switch upper {
	case "SELECT", "FROM", "WHERE", "LIMIT", "OFFSET", "ORDER", "BY", "ASC", "DESC", "NULL", "TRUE", "FALSE", "GROUP", "HAVING":
		return true
	}
	return false
}

func isOperator(lit string) bool {
	upper := strings.ToUpper(lit)
	switch upper {
	case "AND", "OR", "IN", "LIKE", "NOT", "=", "!=", "<", ">", "<=", ">=":
		return true
	}
	return false
}