package jsonpath

import (
	"errors"
	"fmt"
	"github.com/speakeasy-api/jsonpath/pkg/jsonpath/config"
	"github.com/speakeasy-api/jsonpath/pkg/jsonpath/token"
	"strconv"
	"strings"
)

const MaxSafeFloat int64 = 9007199254740991

type mode int

const (
	modeNormal mode = iota
	modeSingular
)

// JSONPath represents a JSONPath parser.
type JSONPath struct {
	tokenizer *token.Tokenizer
	tokens    []token.TokenInfo
	ast       jsonPathAST
	current   int
	mode      []mode
	config    config.Config
}

// newParserPrivate creates a new JSONPath with the given tokens.
func newParserPrivate(tokenizer *token.Tokenizer, tokens []token.TokenInfo, opts ...config.Option) *JSONPath {
	return &JSONPath{tokenizer, tokens, jsonPathAST{}, 0, []mode{modeNormal}, config.New(opts...)}
}

// parse parses the JSONPath tokens and returns the root node of the AST.
//
//	jsonpath-query      = root-identifier segments
func (p *JSONPath) parse() error {
	if len(p.tokens) == 0 {
		return fmt.Errorf("empty JSONPath expression")
	}

	if p.tokens[p.current].Token != token.ROOT {
		return p.parseFailure(&p.tokens[p.current], "expected '$'")
	}
	p.current++

	for p.current < len(p.tokens) {
		segment, err := p.parseSegment()
		if err != nil {
			return err
		}
		p.ast.segments = append(p.ast.segments, segment)
	}
	return nil
}

func (p *JSONPath) parseFailure(target *token.TokenInfo, msg string) error {
	return errors.New(p.tokenizer.ErrorString(target, msg))
}

// peek returns true if the upcoming token matches the given token type.
func (p *JSONPath) peek(token token.Token) bool {
	return p.current+1 < len(p.tokens) && p.tokens[p.current+1].Token == token
}

// peek returns true if the upcoming token matches the given token type.
func (p *JSONPath) next(token token.Token) bool {
	return p.current < len(p.tokens) && p.tokens[p.current].Token == token
}

// expect consumes the current token if it matches the given token type.
func (p *JSONPath) expect(token token.Token) bool {
	if p.peek(token) {
		p.current++
		return true
	}
	return false
}

// isComparisonOperator returns true if the given token is a comparison operator.
func (p *JSONPath) isComparisonOperator(tok token.Token) bool {
	return tok == token.EQ || tok == token.NE || tok == token.GT || tok == token.GE || tok == token.LT || tok == token.LE
}

func (p *JSONPath) parseSegment() (*segment, error) {
	currentToken := p.tokens[p.current]
	if currentToken.Token == token.RECURSIVE {
		if p.mode[len(p.mode)-1] == modeSingular {
			return nil, p.parseFailure(&p.tokens[p.current], "unexpected recursive descent in singular query")
		}
		p.current++
		child, err := p.parseInnerSegment()
		if err != nil {
			return nil, err
		}
		return &segment{kind: segmentKindDescendant, descendant: child}, nil
	} else if currentToken.Token == token.CHILD || currentToken.Token == token.BRACKET_LEFT {
		if currentToken.Token == token.CHILD {
			p.current++
		}
		child, err := p.parseInnerSegment()
		if err != nil {
			return nil, err
		}
		return &segment{kind: segmentKindChild, child: child}, nil
	} else if p.config.PropertyNameEnabled() && currentToken.Token == token.PROPERTY_NAME {
		p.current++
		return &segment{kind: segmentKindProperyName}, nil
	}
	return nil, p.parseFailure(&currentToken, "unexpected token when parsing segment")
}

func (p *JSONPath) parseInnerSegment() (retValue *innerSegment, err error) {
	defer func() {
		if p.mode[len(p.mode)-1] == modeSingular && retValue != nil {
			if len(retValue.selectors) > 1 {
				retValue = nil
				err = p.parseFailure(&p.tokens[p.current], "unexpected multiple selectors in singular query")
				return
			} else if retValue.kind == segmentDotWildcard {
				retValue = nil
				err = p.parseFailure(&p.tokens[p.current], "unexpected wildcard in singular query")
				return
			}
		}
	}()
	// .*
	// .STRING
	// []
	if p.current >= len(p.tokens) {
		return nil, p.parseFailure(nil, "unexpected end of input")
	}
	firstToken := p.tokens[p.current]
	if firstToken.Token == token.WILDCARD {
		p.current += 1
		return &innerSegment{segmentDotWildcard, "", nil}, nil
	} else if firstToken.Token == token.STRING {
		dotName := p.tokens[p.current].Literal
		p.current += 1
		return &innerSegment{segmentDotMemberName, dotName, nil}, nil
	} else if firstToken.Token == token.BRACKET_LEFT {
		prior := p.current
		p.current += 1
		selectors := []*selector{}
		for p.current < len(p.tokens) {
			innerSelector, err := p.parseSelector()
			if err != nil {
				p.current = prior
				return nil, err
			}
			selectors = append(selectors, innerSelector)
			if len(p.tokens) <= p.current {
				return nil, p.parseFailure(&p.tokens[p.current-1], "unexpected end of input")
			}
			if p.tokens[p.current].Token == token.BRACKET_RIGHT {
				break
			} else if p.tokens[p.current].Token == token.COMMA {
				p.current++
			}
		}
		if p.tokens[p.current].Token != token.BRACKET_RIGHT {
			prior = p.current
			return nil, p.parseFailure(&p.tokens[p.current], "expected ']'")
		}
		p.current += 1
		return &innerSegment{kind: segmentLongHand, dotName: "", selectors: selectors}, nil
	}
	return nil, p.parseFailure(&firstToken, "unexpected token when parsing inner segment")
}

func (p *JSONPath) parseSelector() (retSelector *selector, err error) {
	//selector            = name-selector /
	//                      wildcard-selector /
	//                      slice-selector /
	//                      index-selector /
	//                      filter-selector
	initial := p.current
	defer func() {
		if p.mode[len(p.mode)-1] == modeSingular && retSelector != nil {
			if retSelector.kind == selectorSubKindWildcard {
				err = p.parseFailure(&p.tokens[initial], "unexpected wildcard in singular query")
				retSelector = nil
			} else if retSelector.kind == selectorSubKindArraySlice {
				err = p.parseFailure(&p.tokens[initial], "unexpected slice in singular query")
				retSelector = nil
			}
		}
	}()

	//    name-selector       = string-literal
	if p.tokens[p.current].Token == token.STRING_LITERAL {
		name := p.tokens[p.current].Literal
		p.current++
		return &selector{kind: selectorSubKindName, name: name}, nil
		//    wildcard-selector   = "*"
	} else if p.tokens[p.current].Token == token.WILDCARD {
		p.current++
		return &selector{kind: selectorSubKindWildcard}, nil
	} else if p.tokens[p.current].Token == token.INTEGER {
		// peek ahead to see if it's a slice
		if p.peek(token.ARRAY_SLICE) {
			slice, err := p.parseSliceSelector()
			if err != nil {
				return nil, err
			}
			return &selector{kind: selectorSubKindArraySlice, slice: slice}, nil
		}
		// peek ahead to see if we close the array index properly
		if !p.peek(token.BRACKET_RIGHT) && !p.peek(token.COMMA) {
			return nil, p.parseFailure(&p.tokens[p.current], "expected ']' or ','")
		}
		// else it's an index
		lit := p.tokens[p.current].Literal
		// make sure it's not -0
		if lit == "-0" {
			return nil, p.parseFailure(&p.tokens[p.current], "-0 unexpected")
		}
		// make sure lit is an integer
		i, err := strconv.ParseInt(lit, 10, 64)
		if err != nil {
			return nil, p.parseFailure(&p.tokens[p.current], "expected an integer")
		}
		err = p.checkSafeInteger(i, lit)
		if err != nil {
			return nil, err
		}

		p.current++

		return &selector{kind: selectorSubKindArrayIndex, index: i}, nil
	} else if p.tokens[p.current].Token == token.ARRAY_SLICE {
		slice, err := p.parseSliceSelector()
		if err != nil {
			return nil, err
		}
		return &selector{kind: selectorSubKindArraySlice, slice: slice}, nil
	} else if p.tokens[p.current].Token == token.FILTER {
		return p.parseFilterSelector()
	}

	return nil, p.parseFailure(&p.tokens[p.current], "unexpected token when parsing selector")
}

func (p *JSONPath) parseSliceSelector() (*slice, error) {
	// slice-selector = [start S] ":" S [end S] [":" [S step]]
	var start, end, step *int64

	// parse the start index
	if p.tokens[p.current].Token == token.INTEGER {
		literal := p.tokens[p.current].Literal
		i, err := strconv.ParseInt(literal, 10, 64)
		if err != nil {
			return nil, p.parseFailure(&p.tokens[p.current], "expected an integer")
		}
		err = p.checkSafeInteger(i, literal)
		if err != nil {
			return nil, err
		}

		start = &i
		p.current += 1
	}

	// Expect a colon
	if p.tokens[p.current].Token != token.ARRAY_SLICE {
		return nil, p.parseFailure(&p.tokens[p.current], "expected ':'")
	}
	p.current++

	// parse the end index
	if p.tokens[p.current].Token == token.INTEGER {
		literal := p.tokens[p.current].Literal
		i, err := strconv.ParseInt(literal, 10, 64)
		if err != nil {
			return nil, p.parseFailure(&p.tokens[p.current], "expected an integer")
		}
		err = p.checkSafeInteger(i, literal)
		if err != nil {
			return nil, err
		}

		end = &i
		p.current++
	}

	// Check for an optional second colon and step value
	if p.tokens[p.current].Token == token.ARRAY_SLICE {
		p.current++
		if p.tokens[p.current].Token == token.INTEGER {
			literal := p.tokens[p.current].Literal
			i, err := strconv.ParseInt(literal, 10, 64)
			if err != nil {
				return nil, p.parseFailure(&p.tokens[p.current], "expected an integer")
			}
			err = p.checkSafeInteger(i, literal)
			if err != nil {
				return nil, err
			}

			step = &i
			p.current++
		}
	}
	if p.tokens[p.current].Token != token.BRACKET_RIGHT {
		return nil, p.parseFailure(&p.tokens[p.current], "expected ']'")
	}

	return &slice{start: start, end: end, step: step}, nil
}

func (p *JSONPath) checkSafeInteger(i int64, literal string) error {
	if i > MaxSafeFloat || i < -MaxSafeFloat {
		return p.parseFailure(&p.tokens[p.current], "outside bounds for safe integers")
	}
	if literal == "-0" {
		return p.parseFailure(&p.tokens[p.current], "-0 unexpected")
	}
	return nil
}

func (p *JSONPath) parseFilterSelector() (*selector, error) {

	if p.tokens[p.current].Token != token.FILTER {
		return nil, p.parseFailure(&p.tokens[p.current], "expected '?'")
	}
	p.current++

	expr, err := p.parseLogicalOrExpr()
	if err != nil {
		return nil, err
	}

	return &selector{kind: selectorSubKindFilter, filter: &filterSelector{expr}}, nil
}

func (p *JSONPath) parseLogicalOrExpr() (*logicalOrExpr, error) {
	var expr logicalOrExpr

	for {
		andExpr, err := p.parseLogicalAndExpr()
		if err != nil {
			return nil, err
		}
		expr.expressions = append(expr.expressions, andExpr)

		if !p.next(token.OR) {
			break
		}
		p.current++
	}

	return &expr, nil
}

func (p *JSONPath) parseLogicalAndExpr() (*logicalAndExpr, error) {
	var expr logicalAndExpr

	for {
		basicExpr, err := p.parseBasicExpr()
		if err != nil {
			return nil, err
		}
		expr.expressions = append(expr.expressions, basicExpr)

		if !p.next(token.AND) {
			break
		}
		p.current++
	}

	return &expr, nil
}

func (p *JSONPath) parseBasicExpr() (*basicExpr, error) {
	//basic-expr          = paren-expr /
	//	                    comparison-expr /
	//                      test-expr

	switch p.tokens[p.current].Token {
	case token.NOT:
		p.current++
		expr, err := p.parseLogicalOrExpr()
		if err != nil {
			return nil, err
		}
		// Inspect if the expr is topped by a parenExpr -- if so we can simplify
		if len(expr.expressions) == 1 && len(expr.expressions[0].expressions) == 1 && expr.expressions[0].expressions[0].parenExpr != nil {
			child := expr.expressions[0].expressions[0].parenExpr
			child.not = !child.not
			return &basicExpr{parenExpr: child}, nil
		}
		return &basicExpr{parenExpr: &parenExpr{not: true, expr: expr}}, nil
	case token.PAREN_LEFT:
		p.current++
		expr, err := p.parseLogicalOrExpr()
		if err != nil {
			return nil, err
		}
		if p.tokens[p.current].Token != token.PAREN_RIGHT {
			return nil, p.parseFailure(&p.tokens[p.current], "expected ')'")
		}
		p.current++
		return &basicExpr{parenExpr: &parenExpr{not: false, expr: expr}}, nil
	}
	prevCurrent := p.current
	comparisonExpr, comparisonErr := p.parseComparisonExpr()
	if comparisonErr == nil {
		return &basicExpr{comparisonExpr: comparisonExpr}, nil
	}
	p.current = prevCurrent
	testExpr, testErr := p.parseTestExpr()
	if testErr == nil {
		return &basicExpr{testExpr: testExpr}, nil
	}
	p.current = prevCurrent
	return nil, p.parseFailure(&p.tokens[p.current], fmt.Sprintf("could not parse query: expected either testExpr [err: %s] or comparisonExpr: [err: %s]", testErr.Error(), comparisonErr.Error()))
}

func (p *JSONPath) parseComparisonExpr() (*comparisonExpr, error) {
	left, err := p.parseComparable()
	if err != nil {
		return nil, err
	}

	if !p.isComparisonOperator(p.tokens[p.current].Token) {
		return nil, p.parseFailure(&p.tokens[p.current], "expected comparison operator")
	}
	operator := p.tokens[p.current].Token
	var op comparisonOperator
	switch operator {
	case token.EQ:
		op = equalTo
	case token.NE:
		op = notEqualTo
	case token.LT:
		op = lessThan
	case token.LE:
		op = lessThanEqualTo
	case token.GT:
		op = greaterThan
	case token.GE:
		op = greaterThanEqualTo
	default:
		return nil, p.parseFailure(&p.tokens[p.current], "expected comparison operator")
	}
	p.current++

	right, err := p.parseComparable()
	if err != nil {
		return nil, err
	}

	return &comparisonExpr{left: left, op: op, right: right}, nil
}

func (p *JSONPath) parseComparable() (*comparable, error) {
	//	comparable = literal /
	//	singular-query / ; singular query value
	//	function-expr    ; ValueType
	if literal, err := p.parseLiteral(); err == nil {
		return &comparable{literal: literal}, nil
	}
	if funcExpr, err := p.parseFunctionExpr(); err == nil {
		if funcExpr.funcType == functionTypeMatch {
			return nil, p.parseFailure(&p.tokens[p.current], "match result cannot be compared")
		} else if funcExpr.funcType == functionTypeSearch {
			return nil, p.parseFailure(&p.tokens[p.current], "search result cannot be compared")
		}
		return &comparable{functionExpr: funcExpr}, nil
	}
	switch p.tokens[p.current].Token {
	case token.ROOT:
		p.current++
		query, err := p.parseSingleQuery()
		if err != nil {
			return nil, err
		}
		return &comparable{singularQuery: &singularQuery{absQuery: &absQuery{segments: query.segments}}}, nil
	case token.CURRENT:
		p.current++
		query, err := p.parseSingleQuery()
		if err != nil {
			return nil, err
		}
		return &comparable{singularQuery: &singularQuery{relQuery: &relQuery{segments: query.segments}}}, nil
	default:
		return nil, p.parseFailure(&p.tokens[p.current], "expected literal or query")
	}
}

func (p *JSONPath) parseQuery() (*jsonPathAST, error) {
	var query jsonPathAST
	p.mode = append(p.mode, modeNormal)

	for p.current < len(p.tokens) {
		prior := p.current
		segment, err := p.parseSegment()
		if err != nil {
			p.current = prior
			break
		}
		query.segments = append(query.segments, segment)
	}
	p.mode = p.mode[:len(p.mode)-1]
	return &query, nil
}

func (p *JSONPath) parseTestExpr() (*testExpr, error) {
	//test-expr           = [logical-not-op S]
	//                  (filter-query / ; existence/non-existence
	//                   function-expr) ; LogicalType or NodesType
	//filter-query        = rel-query / jsonpath-query
	//rel-query           = current-node-identifier segments
	//current-node-identifier = "@"
	not := false
	if p.tokens[p.current].Token == token.NOT {
		not = true
		p.current++
	}
	switch p.tokens[p.current].Token {
	case token.CURRENT:
		p.current++
		query, err := p.parseQuery()
		if err != nil {
			return nil, err
		}
		return &testExpr{filterQuery: &filterQuery{relQuery: &relQuery{segments: query.segments}}, not: not}, nil
	case token.ROOT:
		p.current++
		query, err := p.parseQuery()
		if err != nil {
			return nil, err
		}
		return &testExpr{filterQuery: &filterQuery{jsonPathQuery: &jsonPathAST{segments: query.segments}}, not: not}, nil
	default:
		funcExpr, err := p.parseFunctionExpr()
		if err != nil {
			return nil, err
		}
		if funcExpr.funcType == functionTypeCount {
			return nil, p.parseFailure(&p.tokens[p.current], "count function must be compared")
		}
		if funcExpr.funcType == functionTypeLength {
			return nil, p.parseFailure(&p.tokens[p.current], "length function must be compared")
		}
		if funcExpr.funcType == functionTypeValue {
			return nil, p.parseFailure(&p.tokens[p.current], "length function must be compared")
		}
		return &testExpr{functionExpr: funcExpr, not: not}, nil
	}

	return nil, p.parseFailure(&p.tokens[p.current], "unexpected token when parsing test expression")
}

func (p *JSONPath) parseFunctionExpr() (*functionExpr, error) {
	functionName := p.tokens[p.current].Literal
	if p.current+1 >= len(p.tokens) || p.tokens[p.current+1].Token != token.PAREN_LEFT {
		return nil, p.parseFailure(&p.tokens[p.current], "expected '(' after function")
	}
	p.current += 2
	args := []*functionArgument{}
	switch functionTypeMap[functionName] {
	case functionTypeLength:
		arg, err := p.parseFunctionArgument(true)
		if err != nil {
			return nil, err
		}
		args = append(args, arg)
	case functionTypeCount:
		arg, err := p.parseFunctionArgument(false)
		if err != nil {
			return nil, err
		}
		if arg.literal != nil && arg.literal.node == nil {
			return nil, p.parseFailure(&p.tokens[p.current], "count function only supports containers")
		}
		args = append(args, arg)
	case functionTypeValue:
		arg, err := p.parseFunctionArgument(false)
		if err != nil {
			return nil, err
		}
		args = append(args, arg)
	case functionTypeMatch:
		fallthrough
	case functionTypeSearch:
		arg, err := p.parseFunctionArgument(false)
		if err != nil {
			return nil, err
		}
		args = append(args, arg)
		if p.tokens[p.current].Token != token.COMMA {
			return nil, p.parseFailure(&p.tokens[p.current], "expected ','")
		}
		p.current++
		arg, err = p.parseFunctionArgument(false)
		if err != nil {
			return nil, err
		}
		args = append(args, arg)
	}
	if p.tokens[p.current].Token != token.PAREN_RIGHT {
		return nil, p.parseFailure(&p.tokens[p.current], "expected ')'")
	}
	p.current++
	return &functionExpr{funcType: functionTypeMap[functionName], args: args}, nil
}

func (p *JSONPath) parseSingleQuery() (*jsonPathAST, error) {
	var query jsonPathAST
	for p.current < len(p.tokens) {
		try := p.current
		p.mode = append(p.mode, modeSingular)
		segment, err := p.parseSegment()
		if err != nil {
			// rollback
			p.mode = p.mode[:len(p.mode)-1]
			p.current = try
			break
		}
		p.mode = p.mode[:len(p.mode)-1]
		query.segments = append(query.segments, segment)
	}
	//if len(query.segments) == 0 {
	//	return nil, p.parseFailure(p.tokens[p.current], "expected at least one segment")
	//}
	return &query, nil
}

func (p *JSONPath) parseFunctionArgument(single bool) (*functionArgument, error) {
	//function-argument   = literal /
	//	filter-query / ; (includes singular-query)
	//  logical-expr /
	//	function-expr

	if lit, err := p.parseLiteral(); err == nil {
		return &functionArgument{literal: lit}, nil
	}
	switch p.tokens[p.current].Token {
	case token.CURRENT:
		p.current++
		var query *jsonPathAST
		var err error
		if single {
			query, err = p.parseSingleQuery()
		} else {
			query, err = p.parseQuery()
		}
		if err != nil {
			return nil, err
		}
		return &functionArgument{filterQuery: &filterQuery{relQuery: &relQuery{segments: query.segments}}}, nil
	case token.ROOT:
		p.current++
		var query *jsonPathAST
		var err error
		if single {
			query, err = p.parseSingleQuery()
		} else {
			query, err = p.parseQuery()
		}
		if err != nil {
			return nil, err
		}
		return &functionArgument{filterQuery: &filterQuery{jsonPathQuery: &jsonPathAST{segments: query.segments}}}, nil
	}
	if expr, err := p.parseLogicalOrExpr(); err == nil {
		return &functionArgument{logicalExpr: expr}, nil
	}
	if funcExpr, err := p.parseFunctionExpr(); err == nil {
		return &functionArgument{functionExpr: funcExpr}, nil
	}

	return nil, p.parseFailure(&p.tokens[p.current], "unexpected token for function argument")
}

func (p *JSONPath) parseLiteral() (*literal, error) {
	switch p.tokens[p.current].Token {
	case token.STRING_LITERAL:
		lit := p.tokens[p.current].Literal
		p.current++
		return &literal{string: &lit}, nil
	case token.INTEGER:
		lit := p.tokens[p.current].Literal
		p.current++
		i, err := strconv.Atoi(lit)
		if err != nil {
			return nil, p.parseFailure(&p.tokens[p.current], "expected integer")
		}
		return &literal{integer: &i}, nil
	case token.FLOAT:
		lit := p.tokens[p.current].Literal
		p.current++
		f, err := strconv.ParseFloat(lit, 64)
		if err != nil {
			return nil, p.parseFailure(&p.tokens[p.current], "expected float")
		}
		return &literal{float64: &f}, nil
	case token.TRUE:
		p.current++
		res := true
		return &literal{bool: &res}, nil
	case token.FALSE:
		p.current++
		res := false
		return &literal{bool: &res}, nil
	case token.NULL:
		p.current++
		res := true
		return &literal{null: &res}, nil
	}
	return nil, p.parseFailure(&p.tokens[p.current], "expected literal")
}

type jsonPathAST struct {
	// "$"
	segments []*segment
}

func (q jsonPathAST) ToString() string {
	b := strings.Builder{}
	b.WriteString("$")
	for _, seg := range q.segments {
		b.WriteString(seg.ToString())
	}
	return b.String()
}
