package router

import (
	"strings"
)

// controlFlow tells the main loop what it should do next.
type controlFlow int

// controlFlow values.
const (
	controlStop  controlFlow = 0
	controlBegin controlFlow = 1
	controlNext  controlFlow = 2
)

// Tree represents a radix tree.
type Tree[T any] struct {
	static      map[string]T
	root        treeNode[T]
	canBeStatic [2048]bool
}

// Add adds a new element to the tree.
func (tree *Tree[T]) Add(path string, data T) {
	if !strings.Contains(path, ":") && !strings.Contains(path, "*") {
		if tree.static == nil {
			tree.static = map[string]T{}
		}

		tree.static[path] = data
		tree.canBeStatic[len(path)] = true
		return
	}

	// Search tree for equal parts until we can no longer proceed
	i := 0
	offset := 0
	node := &tree.root

	for {
	begin:
		switch node.kind {
		case parameter:
			// This only occurs when the same parameter based route is added twice.
			// node: /post/:id|
			// path: /post/:id|
			if i == len(path) {
				node.data = data
				return
			}

			// When we hit a separator, we'll search for a fitting child.
			if path[i] == separator {
				var control controlFlow
				node, offset, control = node.end(path, data, i, offset)

				switch control {
				case controlStop:
					return
				case controlBegin:
					goto begin
				case controlNext:
					goto next
				}
			}

		default:
			if i == len(path) {
				// The path already exists.
				// node: /blog|
				// path: /blog|
				if i-offset == len(node.prefix) {
					node.data = data
					return
				}

				// The path ended but the node prefix is longer.
				// node: /blog|feed
				// path: /blog|
				node.split(i-offset, "", data)
				return
			}

			// The node we just checked is entirely included in our path.
			// node: /|
			// path: /|blog
			if i-offset == len(node.prefix) {
				var control controlFlow
				node, offset, control = node.end(path, data, i, offset)

				switch control {
				case controlStop:
					return
				case controlBegin:
					goto begin
				case controlNext:
					goto next
				}
			}

			// We got a conflict.
			// node: /b|ag
			// path: /b|riefcase
			if path[i] != node.prefix[i-offset] {
				node.split(i-offset, path[i:], data)
				return
			}
		}

	next:
		i++
	}
}

// Lookup finds the data for the given path.
func (tree *Tree[T]) Lookup(path string) (T, []Parameter) {
	var params []Parameter

	data := tree.LookupNoAlloc(path, func(key string, value string) {
		params = append(params, Parameter{key, value})
	})

	return data, params
}

// LookupNoAlloc finds the data for the given path without using any memory allocations.
func (tree *Tree[T]) LookupNoAlloc(path string, addParameter func(key string, value string)) T {
	if tree.canBeStatic[len(path)] {
		handler, found := tree.static[path]

		if found {
			return handler
		}
	}

	var (
		i                  uint
		offset             uint
		lastWildcardOffset uint
		lastWildcard       *treeNode[T]
		empty              T
		node               = &tree.root
	)

begin:
	// Search tree for equal parts until we can no longer proceed
	for {
		// We reached the end.
		if i == uint(len(path)) {
			// node: /blog|
			// path: /blog|
			if i-offset == uint(len(node.prefix)) {
				return node.data
			}

			// node: /blog|feed
			// path: /blog|
			return empty
		}

		// The node we just checked is entirely included in our path.
		// node: /|
		// path: /|blog
		if i-offset == uint(len(node.prefix)) {
			if node.wildcard != nil {
				lastWildcard = node.wildcard
				lastWildcardOffset = i
			}

			char := path[i]

			if char >= node.startIndex && char < node.endIndex {
				index := node.indices[char-node.startIndex]

				if index != 0 {
					node = node.children[index]
					offset = i
					i++
					continue
				}
			}

			// node: /|:id
			// path: /|blog
			if node.parameter != nil {
				node = node.parameter
				offset = i
				i++

				for {
					// We reached the end.
					if i == uint(len(path)) {
						addParameter(node.prefix, path[offset:i])
						return node.data
					}

					// node: /:id|/posts
					// path: /123|/posts
					if path[i] == separator {
						addParameter(node.prefix, path[offset:i])
						index := node.indices[separator-node.startIndex]
						node = node.children[index]
						offset = i
						i++
						goto begin
					}

					i++
				}
			}

			// node: /|*any
			// path: /|image.png
			if node.wildcard != nil {
				addParameter(node.wildcard.prefix, path[i:])
				return node.wildcard.data
			}

			return empty
		}

		// We got a conflict.
		// node: /b|ag
		// path: /b|riefcase
		if path[i] != node.prefix[i-offset] {
			if lastWildcard != nil {
				addParameter(lastWildcard.prefix, path[lastWildcardOffset:])
				return lastWildcard.data
			}

			return empty
		}

		i++
	}
}

// Bind binds all handlers to a new one provided by the callback.
func (tree *Tree[T]) Bind(transform func(T) T) {
	tree.root.each(func(node *treeNode[T]) {
		node.data = transform(node.data)
	})

	for key, value := range tree.static {
		tree.static[key] = transform(value)
	}
}