package main

import (
	"fmt"
	"strings"
	"text/scanner"
)

// Parser parses a string into an AST.
type Parser struct {
	s        scanner.Scanner
}

func NewParser(expr string) *Parser {
	// Preprocess multi-char operators to single chars so the scanner reads them easily
	expr = strings.ReplaceAll(expr, "**", "\x01")
	expr = strings.ReplaceAll(expr, "<<", "\x02")
	expr = strings.ReplaceAll(expr, ">>", "\x03")
	expr = strings.ReplaceAll(expr, "==", "\x04")
	expr = strings.ReplaceAll(expr, "!=", "\x05")
	expr = strings.ReplaceAll(expr, "<=", "\x06")
	expr = strings.ReplaceAll(expr, ">=", "\x07")
	p := &Parser{}
	p.s.Init(strings.NewReader(expr))
	p.s.Mode = scanner.ScanInts | scanner.ScanFloats | scanner.ScanIdents | scanner.ScanStrings
	p.s.Error = func(s *scanner.Scanner, msg string) {}
	return p
}

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

// Parse parses the entire expression.
func (p *Parser) Parse() (Expr, error) {
	expr, err := p.parseExpression()
	if err != nil {
		return nil, err
	}
	p.skipWhitespace()
	if p.s.Peek() != scanner.EOF {
		return nil, fmt.Errorf("unexpected trailing characters")
	}
	return expr, nil
}

func (p *Parser) parseExpression() (Expr, error) {
	return p.parseEquality()
}

func (p *Parser) parseEquality() (Expr, error) {
	res, err := p.parseRelational()
	if err != nil {
		return nil, err
	}
	for {
		p.skipWhitespace()
		ch := p.s.Peek()
		if ch == '\x04' {
			p.s.Next()
			rhs, err := p.parseRelational()
			if err != nil {
				return nil, err
			}
			res = &BinaryExpr{Op: "==", Left: res, Right: rhs}
		} else if ch == '\x05' {
			p.s.Next()
			rhs, err := p.parseRelational()
			if err != nil {
				return nil, err
			}
			res = &BinaryExpr{Op: "!=", Left: res, Right: rhs}
		} else {
			return res, nil
		}
	}
}

func (p *Parser) parseRelational() (Expr, error) {
	res, err := p.parseXor()
	if err != nil {
		return nil, err
	}
	for {
		p.skipWhitespace()
		ch := p.s.Peek()
		if ch == '\x06' {
			p.s.Next()
			rhs, err := p.parseXor()
			if err != nil {
				return nil, err
			}
			res = &BinaryExpr{Op: "<=", Left: res, Right: rhs}
		} else if ch == '\x07' {
			p.s.Next()
			rhs, err := p.parseXor()
			if err != nil {
				return nil, err
			}
			res = &BinaryExpr{Op: ">=", Left: res, Right: rhs}
		} else if ch == '<' {
			p.s.Next()
			rhs, err := p.parseXor()
			if err != nil {
				return nil, err
			}
			res = &BinaryExpr{Op: "<", Left: res, Right: rhs}
		} else if ch == '>' {
			p.s.Next()
			rhs, err := p.parseXor()
			if err != nil {
				return nil, err
			}
			res = &BinaryExpr{Op: ">", Left: res, Right: rhs}
		} else {
			return res, nil
		}
	}
}

func (p *Parser) parseXor() (Expr, error) {
	res, err := p.parseAnd()
	if err != nil {
		return nil, err
	}
	for {
		p.skipWhitespace()
		if p.s.Peek() == '|' || p.s.Peek() == '^' {
			op := string(p.s.Next())
			rhs, err := p.parseAnd()
			if err != nil {
				return nil, err
			}
			res = &BinaryExpr{Op: op, Left: res, Right: rhs}
		} else {
			return res, nil
		}
	}
}

func (p *Parser) parseAnd() (Expr, error) {
	res, err := p.parseShift()
	if err != nil {
		return nil, err
	}
	for {
		p.skipWhitespace()
		if p.s.Peek() == '&' {
			p.s.Next()
			rhs, err := p.parseShift()
			if err != nil {
				return nil, err
			}
			res = &BinaryExpr{Op: "&", Left: res, Right: rhs}
		} else {
			return res, nil
		}
	}
}

func (p *Parser) parseShift() (Expr, error) {
	res, err := p.parseSum()
	if err != nil {
		return nil, err
	}
	for {
		p.skipWhitespace()
		ch := p.s.Peek()
		if ch == '\x02' {
			p.s.Next()
			rhs, err := p.parseSum()
			if err != nil {
				return nil, err
			}
			res = &BinaryExpr{Op: "<<", Left: res, Right: rhs}
		} else if ch == '\x03' {
			p.s.Next()
			rhs, err := p.parseSum()
			if err != nil {
				return nil, err
			}
			res = &BinaryExpr{Op: ">>", Left: res, Right: rhs}
		} else {
			return res, nil
		}
	}
}

func (p *Parser) parseSum() (Expr, error) {
	res, err := p.parseTerm()
	if err != nil {
		return nil, err
	}
	for {
		p.skipWhitespace()
		ch := p.s.Peek()
		if ch == '+' || ch == '-' {
			op := string(p.s.Next())
			rhs, err := p.parseTerm()
			if err != nil {
				return nil, err
			}
			res = &BinaryExpr{Op: op, Left: res, Right: rhs}
		} else {
			return res, nil
		}
	}
}

func (p *Parser) parseTerm() (Expr, error) {
	res, err := p.parsePower()
	if err != nil {
		return nil, err
	}
	for {
		p.skipWhitespace()
		ch := p.s.Peek()
		if ch == '*' || ch == '/' || ch == '%' {
			op := string(p.s.Next())
			rhs, err := p.parsePower()
			if err != nil {
				return nil, err
			}
			res = &BinaryExpr{Op: op, Left: res, Right: rhs}
		} else if (ch >= 'A' && ch <= 'Z') || (ch >= 'a' && ch <= 'z') || ch == '$' || ch == '€' || ch == '£' || ch == '¥' || ch == '(' {
			// Implicit multiplication
			rhs, err := p.parseFactor()
			if err != nil {
				return nil, err
			}
			res = &BinaryExpr{Op: "implicit_mult", Left: res, Right: rhs}
		} else {
			return res, nil
		}
	}
}

func (p *Parser) parsePower() (Expr, error) {
	res, err := p.parseFactor()
	if err != nil {
		return nil, err
	}
	p.skipWhitespace()
	if p.s.Peek() == '\x01' {
		p.s.Next()
		rhs, err := p.parsePower()
		if err != nil {
			return nil, err
		}
		res = &BinaryExpr{Op: "**", Left: res, Right: rhs}
	}
	return res, nil
}

func (p *Parser) parseFactor() (Expr, error) {
	tok := p.s.Scan()
	switch tok {
	case '(':
		res, err := p.parseExpression()
		if err != nil {
			return nil, err
		}
		if p.s.Scan() != ')' {
			return nil, fmt.Errorf("missing )")
		}
		return res, nil
	case '-':
		res, err := p.parseFactor()
		if err != nil {
			return nil, err
		}
		return &UnaryExpr{Op: "-", Expr: res}, nil
	case '~', '!':
		op := string(tok)
		res, err := p.parseFactor()
		if err != nil {
			return nil, err
		}
		return &UnaryExpr{Op: op, Expr: res}, nil
	case '$':
		return &IdentExpr{Name: "$"}, nil
	case '€':
		return &IdentExpr{Name: "€"}, nil
	case '£':
		return &IdentExpr{Name: "£"}, nil
	case '¥':
		return &IdentExpr{Name: "¥"}, nil
	case scanner.String:
		text := p.s.TokenText()
		text = text[1 : len(text)-1]
		return &StringExpr{Value: text}, nil
	case scanner.Int, scanner.Float, scanner.Ident:
		text := p.s.TokenText()
		if tok == scanner.Ident {
			p.skipWhitespace()
			if p.s.Peek() == '(' {
				p.s.Next() // consume '('
				var args []Expr
				for {
					p.skipWhitespace()
					if p.s.Peek() == ')' {
						p.s.Next()
						break
					}
					arg, err := p.parseExpression()
					if err != nil {
						return nil, err
					}
					args = append(args, arg)
					p.skipWhitespace()
					nextChar := p.s.Peek()
					if nextChar == ',' {
						p.s.Next()
						continue
					} else if nextChar == ')' {
						p.s.Next()
						break
					} else {
						return nil, fmt.Errorf("missing ) or comma in function arguments")
					}
				}
				return &CallExpr{Func: text, Args: args}, nil
			}
			return &IdentExpr{Name: text}, nil
		}
		return &NumberExpr{Value: text}, nil
	default:
		txt := p.s.TokenText()
		if tok == scanner.EOF {
			return nil, fmt.Errorf("unexpected end of expression")
		}
		return nil, fmt.Errorf("unexpected %s", txt)
	}
}