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 "*")
}

// JSONNode represents a value in the JSON data with its path
type JSONNode struct {
	Value interface{} // The value found at the path
	Path  string      // The exact JSONPath where the value was found
}

// StepType defines the types of steps in a JSONPath
type StepType int

const (
	RootStep             StepType = iota // $ - The root element
	ChildStep                            // .key - Direct child access
	RecursiveDescentStep                 // ..key - Recursive search for key
	WildcardStep                         // .* - All children of an object
	IndexStep                            // [n] - Array index access (or [*] for all elements)
)

// 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
)

// ParseJSONPath parses a JSONPath string into a sequence of steps
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
}

// readKey extracts a key name from the path
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
}

// readIndex extracts an array index or wildcard from the path
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 with their paths from data at the specified JSONPath
// Each returned JSONNode contains both the value and its exact path in the data structure
func Get(data interface{}, path string) []JSONNode {
	steps, err := ParseJSONPath(path)
	if err != nil {
		log.Println("Error parsing JSONPath:", err)
		return nil
	}

	results := []JSONNode{}
	traverseWithPaths(data, steps, &results, "$")
	return results
}

// Set updates the value at the specified JSONPath in the original data structure.
// It only modifies the first matching node.
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
	setWithPath(data, steps, &success, value, "$", ModifyFirstMode)
	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
	setWithPath(data, steps, &success, value, "$", ModifyAllMode)
	return success
}

// setWithPath modifies values while tracking paths
func setWithPath(node interface{}, steps []JSONStep, success *bool, value interface{}, currentPath string, mode TraversalMode) {
	if node == nil || *success && mode == ModifyFirstMode {
		return
	}

	// Skip root step
	actualSteps := steps
	if len(steps) > 0 && steps[0].Type == RootStep {
		if len(steps) == 1 {
			return // Cannot set root node
		}
		actualSteps = steps[1:]
	}

	// Process the first step
	if len(actualSteps) == 0 {
		return
	}

	step := actualSteps[0]
	remainingSteps := actualSteps[1:]
	isLastStep := len(remainingSteps) == 0

	switch step.Type {
	case ChildStep:
		m, ok := node.(map[string]interface{})
		if !ok {
			return
		}

		childPath := currentPath + "." + step.Key

		if isLastStep {
			// We've reached the target, set the value
			m[step.Key] = value
			*success = true
			return
		}

		// Create intermediate nodes if necessary
		child, exists := m[step.Key]
		if !exists {
			// Create missing intermediate node
			if len(remainingSteps) > 0 && remainingSteps[0].Type == IndexStep {
				child = []interface{}{}
			} else {
				child = map[string]interface{}{}
			}
			m[step.Key] = child
		}

		setWithPath(child, remainingSteps, success, value, childPath, mode)

	case IndexStep:
		arr, ok := node.([]interface{})
		if !ok {
			return
		}

		// Handle wildcard index
		if step.Index == -1 {
			for i, item := range arr {
				itemPath := fmt.Sprintf("%s[%d]", currentPath, i)
				if isLastStep {
					arr[i] = value
					*success = true
					if mode == ModifyFirstMode {
						return
					}
				} else {
					setWithPath(item, remainingSteps, success, value, itemPath, mode)
					if *success && mode == ModifyFirstMode {
						return
					}
				}
			}
			return
		}

		// Handle specific index
		if step.Index >= 0 && step.Index < len(arr) {
			item := arr[step.Index]
			itemPath := fmt.Sprintf("%s[%d]", currentPath, step.Index)
			if isLastStep {
				arr[step.Index] = value
				*success = true
			} else {
				setWithPath(item, remainingSteps, success, value, itemPath, mode)
			}
		}

	case RecursiveDescentStep:
		// For recursive descent, first check direct match at this level
		if m, ok := node.(map[string]interface{}); ok && step.Key != "*" {
			if val, exists := m[step.Key]; exists {
				directPath := currentPath + "." + step.Key
				if isLastStep {
					m[step.Key] = value
					*success = true
					if mode == ModifyFirstMode {
						return
					}
				} else {
					setWithPath(val, remainingSteps, success, value, directPath, mode)
					if *success && mode == ModifyFirstMode {
						return
					}
				}
			}
		}

		// Then continue recursion to all children
		switch n := node.(type) {
		case map[string]interface{}:
			for k, v := range n {
				childPath := currentPath + "." + k
				// Skip keys we've already processed directly
				if step.Key != "*" && k == step.Key {
					continue
				}
				setWithPath(v, steps, success, value, childPath, mode) // Use the same steps
				if *success && mode == ModifyFirstMode {
					return
				}
			}
		case []interface{}:
			for i, v := range n {
				childPath := fmt.Sprintf("%s[%d]", currentPath, i)
				setWithPath(v, steps, success, value, childPath, mode) // Use the same steps
				if *success && mode == ModifyFirstMode {
					return
				}
			}
		}

	case WildcardStep:
		m, ok := node.(map[string]interface{})
		if !ok {
			return
		}

		for k, v := range m {
			childPath := currentPath + "." + k
			if isLastStep {
				m[k] = value
				*success = true
				if mode == ModifyFirstMode {
					return
				}
			} else {
				setWithPath(v, remainingSteps, success, value, childPath, mode)
				if *success && mode == ModifyFirstMode {
					return
				}
			}
		}
	}
}

// traverseWithPaths tracks both nodes and their paths during traversal
func traverseWithPaths(node interface{}, steps []JSONStep, results *[]JSONNode, currentPath string) {
	if len(steps) == 0 || node == nil {
		return
	}

	// Skip root step
	actualSteps := steps
	if steps[0].Type == RootStep {
		if len(steps) == 1 {
			*results = append(*results, JSONNode{Value: node, Path: currentPath})
			return
		}
		actualSteps = steps[1:]
	}

	// Process the first step
	step := actualSteps[0]
	remainingSteps := actualSteps[1:]
	isLastStep := len(remainingSteps) == 0

	switch step.Type {
	case ChildStep:
		m, ok := node.(map[string]interface{})
		if !ok {
			return
		}

		child, exists := m[step.Key]
		if !exists {
			return
		}

		childPath := currentPath + "." + step.Key
		if isLastStep {
			*results = append(*results, JSONNode{Value: child, Path: childPath})
		} else {
			traverseWithPaths(child, remainingSteps, results, childPath)
		}

	case IndexStep:
		arr, ok := node.([]interface{})
		if !ok {
			return
		}

		// Handle wildcard index
		if step.Index == -1 {
			for i, item := range arr {
				itemPath := fmt.Sprintf("%s[%d]", currentPath, i)
				if isLastStep {
					*results = append(*results, JSONNode{Value: item, Path: itemPath})
				} else {
					traverseWithPaths(item, remainingSteps, results, itemPath)
				}
			}
			return
		}

		// Handle specific index
		if step.Index >= 0 && step.Index < len(arr) {
			item := arr[step.Index]
			itemPath := fmt.Sprintf("%s[%d]", currentPath, step.Index)
			if isLastStep {
				*results = append(*results, JSONNode{Value: item, Path: itemPath})
			} else {
				traverseWithPaths(item, remainingSteps, results, itemPath)
			}
		}

	case RecursiveDescentStep:
		// For recursive descent, first check direct match at this level
		if m, ok := node.(map[string]interface{}); ok && step.Key != "*" {
			if val, exists := m[step.Key]; exists {
				directPath := currentPath + "." + step.Key
				if isLastStep {
					*results = append(*results, JSONNode{Value: val, Path: directPath})
				} else {
					traverseWithPaths(val, remainingSteps, results, directPath)
				}
			}
		}

		// For wildcard, collect this node
		if step.Key == "*" && isLastStep {
			*results = append(*results, JSONNode{Value: node, Path: currentPath})
		}

		// Then continue recursion to all children
		switch n := node.(type) {
		case map[string]interface{}:
			for k, v := range n {
				childPath := currentPath + "." + k
				traverseWithPaths(v, steps, results, childPath) // Use the same steps
			}
		case []interface{}:
			for i, v := range n {
				childPath := fmt.Sprintf("%s[%d]", currentPath, i)
				traverseWithPaths(v, steps, results, childPath) // Use the same steps
			}
		}

	case WildcardStep:
		m, ok := node.(map[string]interface{})
		if !ok {
			return
		}

		for k, v := range m {
			childPath := currentPath + "." + k
			if isLastStep {
				*results = append(*results, JSONNode{Value: v, Path: childPath})
			} else {
				traverseWithPaths(v, remainingSteps, results, childPath)
			}
		}
	}
}