package router

import (
	"errors"
	"net/http"
	"path"
	"strings"
)

var Wildcard = "{}"
var WildcardHeader = "__wildcard_value__"

type tree struct {
	next    map[string]*tree
	handler http.HandlerFunc
}

func newTree() *tree {
	return &tree{
		next:    make(map[string]*tree),
		handler: nil,
	}
}

func (t *tree) Lookup(path string) http.HandlerFunc {
	if path == "/" || path == "" {
		if t.handler != nil {
			return t.handler
		}
		if n, ok := t.next[Wildcard+Wildcard]; t.handler == nil && ok {
			foo := n.handler
			return foo
		}
		return nil
	}
	key, following := nextPathSegment(path)
	n, ok := t.next[key]
	if ok {
		n2 := n.Lookup(following)
		if n2 != nil {
			return n2
		}
	}
	if n, ok := t.next[Wildcard]; ok {
		foo := n.Lookup(following)
		if foo != nil {
			return func(w http.ResponseWriter, r *http.Request) {
				r.Header.Add(WildcardHeader, key)
				foo(w, r)
			}
		}
	} else if n, ok := t.next[Wildcard+Wildcard]; ok {
		foo := n.handler
		if foo != nil {
			return func(w http.ResponseWriter, r *http.Request) {
				r.Header.Add(WildcardHeader, key)
				foo(w, r)
			}
		}
	}
	return nil
}

func (t *tree) Insert(path string, foo http.HandlerFunc) error {
	if path == "/" {
		if t.handler != nil {
			return errors.New("occupied path")
		}
		t.handler = foo
		return nil
	}
	key, following := nextPathSegment(path)
	_, ok := t.next[key]
	if !ok {
		t.next[key] = newTree()
	}
	return t.next[key].Insert(following, foo)
}

func nextPathSegment(p string) (string, string) {
	p = path.Clean("/" + p)
	i := strings.Index(p[1:], "/") + 1
	if i <= 0 {
		return p[1:], "/"
	}
	return p[1:i], p[i:]
}