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
)

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
}

func EvaluateJSONPath(data interface{}, steps []JSONStep) []interface{} {
	current := []interface{}{data}

	for _, step := range steps {
		var next []interface{}
		for _, node := range current {
			next = append(next, evalStep(node, step)...)
		}
		current = next
	}

	return current
}

func evalStep(node interface{}, step JSONStep) []interface{} {
	switch step.Type {
	case ChildStep:
		return evalChild(node, step.Key)
	case RecursiveDescentStep:
		return evalRecursiveDescent(node, step.Key)
	case WildcardStep:
		return evalWildcard(node)
	case IndexStep:
		return evalIndex(node, step.Index)
	case RootStep:
		return []interface{}{node}
	default:
		log.Println("Unknown step type:", step.Type)
		return nil
	}
}

func evalChild(node interface{}, key string) []interface{} {
	if m, ok := node.(map[string]interface{}); ok {
		if val, exists := m[key]; exists {
			return []interface{}{val}
		}
	}
	return nil
}

func evalRecursiveDescent(node interface{}, targetKey string) []interface{} {
	results := []interface{}{}
	queue := []interface{}{node}

	for len(queue) > 0 {
		current := queue[0]
		queue = queue[1:]

		if targetKey == "*" {
			results = append(results, current)
		} else if m, ok := current.(map[string]interface{}); ok {
			if val, exists := m[targetKey]; exists {
				results = append(results, val)
			}
		}

		if m, ok := current.(map[string]interface{}); ok {
			for _, v := range m {
				queue = append(queue, v)
			}
		} else if arr, ok := current.([]interface{}); ok {
			queue = append(queue, arr...)
		}
	}

	return results
}

func evalWildcard(node interface{}) []interface{} {
	if m, ok := node.(map[string]interface{}); ok {
		results := make([]interface{}, 0, len(m))
		for _, v := range m {
			results = append(results, v)
		}
		return results
	}
	return nil
}

func evalIndex(node interface{}, index int) []interface{} {
	if arr, ok := node.([]interface{}); ok {
		if index == -1 { // Wildcard [*]
			return arr
		}
		if index >= 0 && index < len(arr) {
			return []interface{}{arr[index]}
		}
	}
	return nil
}