package jsonpath

import (
	"fmt"
	"log"
	"strconv"
)

// JSONStep represents a single step in a JSONPath query
type JSONStep struct {
	Type  StepType
	Key   string // For Child/RecursiveDescent
	Index int    // For Index (use -1 for wildcard "*")
}

type StepType int

const (
	RootStep StepType = iota
	ChildStep
	RecursiveDescentStep
	WildcardStep
	IndexStep
)

// TraversalMode determines how the traversal behaves
type TraversalMode int

const (
	CollectMode     TraversalMode = iota // Just collect matched nodes
	ModifyFirstMode                      // Modify first matching node
	ModifyAllMode                        // Modify all matching nodes
)

func ParseJSONPath(path string) ([]JSONStep, error) {
	if len(path) == 0 || path[0] != '$' {
		return nil, fmt.Errorf("path must start with $")
	}

	steps := []JSONStep{}
	i := 0

	for i < len(path) {
		switch path[i] {
		case '$':
			steps = append(steps, JSONStep{Type: RootStep})
			i++
		case '.':
			i++
			if i < len(path) && path[i] == '.' {
				// Recursive descent
				i++
				key, nextPos := readKey(path, i)
				steps = append(steps, JSONStep{Type: RecursiveDescentStep, Key: key})
				i = nextPos
			} else {
				// Child step or wildcard
				key, nextPos := readKey(path, i)
				if key == "*" {
					steps = append(steps, JSONStep{Type: WildcardStep})
				} else {
					steps = append(steps, JSONStep{Type: ChildStep, Key: key})
				}
				i = nextPos
			}
		case '[':
			// Index step
			i++
			indexStr, nextPos := readIndex(path, i)
			if indexStr == "*" {
				steps = append(steps, JSONStep{Type: IndexStep, Index: -1})
			} else {
				index, err := strconv.Atoi(indexStr)
				if err != nil {
					return nil, fmt.Errorf("invalid index: %s", indexStr)
				}
				steps = append(steps, JSONStep{Type: IndexStep, Index: index})
			}
			i = nextPos + 1 // Skip closing ]
		default:
			return nil, fmt.Errorf("unexpected character: %c", path[i])
		}
	}

	return steps, nil
}

func readKey(path string, start int) (string, int) {
	i := start
	for ; i < len(path); i++ {
		if path[i] == '.' || path[i] == '[' {
			break
		}
	}
	return path[start:i], i
}

func readIndex(path string, start int) (string, int) {
	i := start
	for ; i < len(path); i++ {
		if path[i] == ']' {
			break
		}
	}
	return path[start:i], i
}

// Get retrieves values from data at the specified JSONPath
func Get(data interface{}, path string) []interface{} {
	steps, err := ParseJSONPath(path)
	if err != nil {
		log.Println("Error parsing JSONPath:", err)
		return nil
	}
	results := []interface{}{}
	traverse(data, steps, CollectMode, nil, &results)
	return results
}

// Set updates the value at the specified JSONPath in the original data structure.
func Set(data interface{}, path string, value interface{}) bool {
	steps, err := ParseJSONPath(path)
	if err != nil {
		log.Println("Error parsing JSONPath:", err)
		return false
	}

	if len(steps) <= 1 {
		log.Println("Cannot set root node")
		return false
	}

	success := false
	traverse(data, steps, ModifyFirstMode, value, &success)
	return success
}

// SetAll updates all matching values at the specified JSONPath.
func SetAll(data interface{}, path string, value interface{}) bool {
	steps, err := ParseJSONPath(path)
	if err != nil {
		log.Println("Error parsing JSONPath:", err)
		return false
	}

	if len(steps) <= 1 {
		log.Println("Cannot set root node")
		return false
	}

	success := false
	traverse(data, steps, ModifyAllMode, value, &success)
	return success
}

// traverse is the main entry point for JSONPath traversal
func traverse(data interface{}, steps []JSONStep, mode TraversalMode, value interface{}, result interface{}) {
	if len(steps) == 0 || data == nil {
		return
	}

	// Skip root step
	actualSteps := steps
	if steps[0].Type == RootStep {
		if len(steps) == 1 {
			if mode == CollectMode {
				results := result.(*[]interface{})
				*results = append(*results, data)
			}
			return
		}
		actualSteps = steps[1:]
	}

	// Start recursion
	traverseSteps(data, actualSteps, mode, value, result)
}

// traverseSteps handles all path traversal with a unified algorithm
func traverseSteps(node interface{}, steps []JSONStep, mode TraversalMode, value interface{}, result interface{}) {
	if len(steps) == 0 {
		// We've reached a terminal node
		if mode == CollectMode {
			results := result.(*[]interface{})
			*results = append(*results, node)
		}
		return
	}

	// Extract current step info
	step := steps[0]
	isLastStep := len(steps) == 1
	nextSteps := steps[1:]

	// For modify modes, get the success pointer
	var successPtr *bool
	if mode != CollectMode {
		successPtr = result.(*bool)
	}

	// Helper function to check if we should stop recursion
	shouldStop := func() bool {
		return mode == ModifyFirstMode && *successPtr
	}

	// Handle each step type
	switch step.Type {
	case ChildStep:
		// Only process maps for child steps
		m, ok := node.(map[string]interface{})
		if !ok {
			return
		}

		child, exists := m[step.Key]

		// Handle modification on last step
		if isLastStep && (mode == ModifyFirstMode || mode == ModifyAllMode) {
			m[step.Key] = value
			*successPtr = true
			return
		}

		// Continue traversal with existing child
		if exists {
			traverseSteps(child, nextSteps, mode, value, result)
			return
		}

		// Create missing intermediate nodes
		if !isLastStep && (mode == ModifyFirstMode || mode == ModifyAllMode) {
			var newNode interface{}
			if len(nextSteps) > 0 && nextSteps[0].Type == IndexStep {
				newNode = []interface{}{}
			} else {
				newNode = map[string]interface{}{}
			}
			m[step.Key] = newNode
			traverseSteps(newNode, nextSteps, mode, value, result)
		}

	case IndexStep:
		// Only process arrays for index steps
		arr, ok := node.([]interface{})
		if !ok {
			return
		}

		// Handle wildcard index
		if step.Index == -1 {
			for i, item := range arr {
				if isLastStep && (mode == ModifyFirstMode || mode == ModifyAllMode) {
					arr[i] = value
					*successPtr = true
					if shouldStop() {
						return
					}
				} else {
					traverseSteps(item, nextSteps, mode, value, result)
					if shouldStop() {
						return
					}
				}
			}
			return
		}

		// Handle specific index
		if step.Index >= 0 && step.Index < len(arr) {
			if isLastStep && (mode == ModifyFirstMode || mode == ModifyAllMode) {
				arr[step.Index] = value
				*successPtr = true
			} else {
				traverseSteps(arr[step.Index], nextSteps, mode, value, result)
			}
		}

	case RecursiveDescentStep:
		// For recursive descent, first collect/modify match at this level if available
		if m, ok := node.(map[string]interface{}); ok && step.Key != "*" {
			if val, exists := m[step.Key]; exists {
				if isLastStep {
					if mode == CollectMode {
						results := result.(*[]interface{})
						*results = append(*results, val)
					} else { // Modify modes
						m[step.Key] = value
						*successPtr = true
						if shouldStop() {
							return
						}
					}
				} else if !isLastStep && mode != CollectMode {
					// Continue with next steps for non-terminal direct matches
					traverseSteps(val, nextSteps, mode, value, result)
					if shouldStop() {
						return
					}
				}
			}
		}

		// For wildcard, collect this node
		if step.Key == "*" && mode == CollectMode {
			results := result.(*[]interface{})
			*results = append(*results, node)
		}

		// Then continue recursion to all children
		switch n := node.(type) {
		case map[string]interface{}:
			for _, v := range n {
				traverseSteps(v, steps, mode, value, result) // Use same steps
				if shouldStop() {
					return
				}
			}
		case []interface{}:
			for _, v := range n {
				traverseSteps(v, steps, mode, value, result) // Use same steps
				if shouldStop() {
					return
				}
			}
		}

	case WildcardStep:
		// Only process maps for wildcard steps
		m, ok := node.(map[string]interface{})
		if !ok {
			return
		}

		// Process all keys
		for k, v := range m {
			if isLastStep && (mode == ModifyFirstMode || mode == ModifyAllMode) {
				m[k] = value
				*successPtr = true
				if shouldStop() {
					return
				}
			} else {
				traverseSteps(v, nextSteps, mode, value, result)
				if shouldStop() {
					return
				}
			}
		}
	}
}