package router

import (
	"fmt"
	"net/http"
)

// Handler is an interface for handling HTTP requests with path parameters.
type Handler interface {
	Serve(params []string)
}

// Middleware wraps a handler with additional functionality.
type Middleware func(Handler) Handler

// node represents a segment in the URL path and its handling logic.
type node struct {
	segment    string  // the path segment this node matches
	handler    Handler // handler for this path, if it's an endpoint
	children   []*node // child nodes for subsequent path segments
	isDynamic  bool    // true for param segments like [id]
	isWildcard bool    // true for catch-all segments like *filepath
	maxParams  uint8   // maximum number of parameters in paths under this node
}

// Router routes HTTP requests by method and path.
// It supports static paths, path parameters, wildcards, and middleware.
type Router struct {
	get        *node
	post       *node
	put        *node
	patch      *node
	delete     *node
	middleware []Middleware // Global middleware
}

// Group represents a route group with a path prefix and shared middleware.
type Group struct {
	router     *Router
	prefix     string
	middleware []Middleware
}

// New creates a new Router instance.
func New() *Router {
	return &Router{
		get:        &node{},
		post:       &node{},
		put:        &node{},
		patch:      &node{},
		delete:     &node{},
		middleware: []Middleware{},
	}
}

// ServeHTTP implements http.Handler interface
func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	handler, params, found := r.Lookup(req.Method, req.URL.Path)
	if !found {
		http.NotFound(w, req)
		return
	}

	// Create an HTTP-specific handler wrapper
	httpHandler := &httpHandler{
		w: w,
		r: req,
		h: func(w http.ResponseWriter, r *http.Request, params []string) {
			handler.Serve(params)
		},
	}

	httpHandler.Serve(params)
}

// httpHandler adapts net/http handlers to the router.
type httpHandler struct {
	w http.ResponseWriter
	r *http.Request
	h func(w http.ResponseWriter, r *http.Request, params []string)
}

// Serve executes the http handler with parameters.
func (h *httpHandler) Serve(params []string) {
	h.h(h.w, h.r, params)
}

// Use adds middleware to the router's global middleware stack.
func (r *Router) Use(middleware ...Middleware) *Router {
	r.middleware = append(r.middleware, middleware...)
	return r
}

// Group creates a new route group with the given path prefix.
func (r *Router) Group(prefix string) *Group {
	return &Group{
		router:     r,
		prefix:     prefix,
		middleware: []Middleware{},
	}
}

// Use adds middleware to the group's middleware stack.
func (g *Group) Use(middleware ...Middleware) *Group {
	g.middleware = append(g.middleware, middleware...)
	return g
}

// Group creates a nested group with an additional prefix.
func (g *Group) Group(prefix string) *Group {
	return &Group{
		router:     g.router,
		prefix:     g.prefix + prefix,
		middleware: append([]Middleware{}, g.middleware...),
	}
}

// applyMiddleware wraps a handler with middleware in reverse order.
func applyMiddleware(handler Handler, middleware []Middleware) Handler {
	h := handler
	for i := len(middleware) - 1; i >= 0; i-- {
		h = middleware[i](h)
	}
	return h
}

// Get registers a handler for GET requests at the given path.
func (r *Router) Get(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	return r.addRoute(r.get, path, &httpHandler{h: httpHandlerFunc}, r.middleware)
}

// Get registers a handler with parameters for GET requests at the given path.
func (r *Router) GetParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	return r.addRoute(r.get, path, &httpHandler{h: handler}, r.middleware)
}

// Post registers a handler for POST requests at the given path.
func (r *Router) Post(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	return r.addRoute(r.post, path, &httpHandler{h: httpHandlerFunc}, r.middleware)
}

// Post registers a handler with parameters for POST requests at the given path.
func (r *Router) PostParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	return r.addRoute(r.post, path, &httpHandler{h: handler}, r.middleware)
}

// Put registers a handler for PUT requests at the given path.
func (r *Router) Put(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	return r.addRoute(r.put, path, &httpHandler{h: httpHandlerFunc}, r.middleware)
}

// Put registers a handler with parameters for PUT requests at the given path.
func (r *Router) PutParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	return r.addRoute(r.put, path, &httpHandler{h: handler}, r.middleware)
}

// Patch registers a handler for PATCH requests at the given path.
func (r *Router) Patch(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	return r.addRoute(r.patch, path, &httpHandler{h: httpHandlerFunc}, r.middleware)
}

// Patch registers a handler with parameters for PATCH requests at the given path.
func (r *Router) PatchParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	return r.addRoute(r.patch, path, &httpHandler{h: handler}, r.middleware)
}

// Delete registers a handler for DELETE requests at the given path.
func (r *Router) Delete(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	return r.addRoute(r.delete, path, &httpHandler{h: httpHandlerFunc}, r.middleware)
}

// Delete registers a handler with parameters for DELETE requests at the given path.
func (r *Router) DeleteParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	return r.addRoute(r.delete, path, &httpHandler{h: handler}, r.middleware)
}

// Group HTTP handler registration methods

// Get registers a handler for GET requests at the given path.
func (g *Group) Get(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	middleware := append([]Middleware{}, g.router.middleware...)
	middleware = append(middleware, g.middleware...)
	return g.router.addRoute(g.router.get, g.prefix+path, &httpHandler{h: httpHandlerFunc}, middleware)
}

// Get registers a handler with parameters for GET requests at the given path.
func (g *Group) GetParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	middleware := append([]Middleware{}, g.router.middleware...)
	middleware = append(middleware, g.middleware...)
	return g.router.addRoute(g.router.get, g.prefix+path, &httpHandler{h: handler}, middleware)
}

// Post registers a handler for POST requests at the given path.
func (g *Group) Post(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	middleware := append([]Middleware{}, g.router.middleware...)
	middleware = append(middleware, g.middleware...)
	return g.router.addRoute(g.router.post, g.prefix+path, &httpHandler{h: httpHandlerFunc}, middleware)
}

// Post registers a handler with parameters for POST requests at the given path.
func (g *Group) PostParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	middleware := append([]Middleware{}, g.router.middleware...)
	middleware = append(middleware, g.middleware...)
	return g.router.addRoute(g.router.post, g.prefix+path, &httpHandler{h: handler}, middleware)
}

// Put registers a handler for PUT requests at the given path.
func (g *Group) Put(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	middleware := append([]Middleware{}, g.router.middleware...)
	middleware = append(middleware, g.middleware...)
	return g.router.addRoute(g.router.put, g.prefix+path, &httpHandler{h: httpHandlerFunc}, middleware)
}

// Put registers a handler with parameters for PUT requests at the given path.
func (g *Group) PutParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	middleware := append([]Middleware{}, g.router.middleware...)
	middleware = append(middleware, g.middleware...)
	return g.router.addRoute(g.router.put, g.prefix+path, &httpHandler{h: handler}, middleware)
}

// Patch registers a handler for PATCH requests at the given path.
func (g *Group) Patch(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	middleware := append([]Middleware{}, g.router.middleware...)
	middleware = append(middleware, g.middleware...)
	return g.router.addRoute(g.router.patch, g.prefix+path, &httpHandler{h: httpHandlerFunc}, middleware)
}

// Patch registers a handler with parameters for PATCH requests at the given path.
func (g *Group) PatchParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	middleware := append([]Middleware{}, g.router.middleware...)
	middleware = append(middleware, g.middleware...)
	return g.router.addRoute(g.router.patch, g.prefix+path, &httpHandler{h: handler}, middleware)
}

// Delete registers a handler for DELETE requests at the given path.
func (g *Group) Delete(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	middleware := append([]Middleware{}, g.router.middleware...)
	middleware = append(middleware, g.middleware...)
	return g.router.addRoute(g.router.delete, g.prefix+path, &httpHandler{h: httpHandlerFunc}, middleware)
}

// Delete registers a handler with parameters for DELETE requests at the given path.
func (g *Group) DeleteParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	middleware := append([]Middleware{}, g.router.middleware...)
	middleware = append(middleware, g.middleware...)
	return g.router.addRoute(g.router.delete, g.prefix+path, &httpHandler{h: handler}, middleware)
}

// WithMiddleware applies specific middleware to the next route registration.
func (r *Router) WithMiddleware(middleware ...Middleware) *MiddlewareRouter {
	return &MiddlewareRouter{
		router:     r,
		middleware: middleware,
	}
}

// WithMiddleware applies specific middleware to the next route registration.
func (g *Group) WithMiddleware(middleware ...Middleware) *MiddlewareGroup {
	return &MiddlewareGroup{
		group:      g,
		middleware: middleware,
	}
}

// MiddlewareRouter handles route registration with specific middleware.
type MiddlewareRouter struct {
	router     *Router
	middleware []Middleware
}

// MiddlewareGroup handles group route registration with specific middleware.
type MiddlewareGroup struct {
	group      *Group
	middleware []Middleware
}

// HTTP method registration functions for MiddlewareRouter

// Get registers a handler for GET requests with specific middleware.
func (mr *MiddlewareRouter) Get(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	middleware := append([]Middleware{}, mr.router.middleware...)
	middleware = append(middleware, mr.middleware...)
	return mr.router.addRoute(mr.router.get, path, &httpHandler{h: httpHandlerFunc}, middleware)
}

// Get registers a handler with parameters for GET requests with specific middleware.
func (mr *MiddlewareRouter) GetParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	middleware := append([]Middleware{}, mr.router.middleware...)
	middleware = append(middleware, mr.middleware...)
	return mr.router.addRoute(mr.router.get, path, &httpHandler{h: handler}, middleware)
}

// Post registers a handler for POST requests with specific middleware.
func (mr *MiddlewareRouter) Post(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	middleware := append([]Middleware{}, mr.router.middleware...)
	middleware = append(middleware, mr.middleware...)
	return mr.router.addRoute(mr.router.post, path, &httpHandler{h: httpHandlerFunc}, middleware)
}

// Post registers a handler with parameters for POST requests with specific middleware.
func (mr *MiddlewareRouter) PostParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	middleware := append([]Middleware{}, mr.router.middleware...)
	middleware = append(middleware, mr.middleware...)
	return mr.router.addRoute(mr.router.post, path, &httpHandler{h: handler}, middleware)
}

// Put registers a handler for PUT requests with specific middleware.
func (mr *MiddlewareRouter) Put(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	middleware := append([]Middleware{}, mr.router.middleware...)
	middleware = append(middleware, mr.middleware...)
	return mr.router.addRoute(mr.router.put, path, &httpHandler{h: httpHandlerFunc}, middleware)
}

// Put registers a handler with parameters for PUT requests with specific middleware.
func (mr *MiddlewareRouter) PutParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	middleware := append([]Middleware{}, mr.router.middleware...)
	middleware = append(middleware, mr.middleware...)
	return mr.router.addRoute(mr.router.put, path, &httpHandler{h: handler}, middleware)
}

// Patch registers a handler for PATCH requests with specific middleware.
func (mr *MiddlewareRouter) Patch(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	middleware := append([]Middleware{}, mr.router.middleware...)
	middleware = append(middleware, mr.middleware...)
	return mr.router.addRoute(mr.router.patch, path, &httpHandler{h: httpHandlerFunc}, middleware)
}

// Patch registers a handler with parameters for PATCH requests with specific middleware.
func (mr *MiddlewareRouter) PatchParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	middleware := append([]Middleware{}, mr.router.middleware...)
	middleware = append(middleware, mr.middleware...)
	return mr.router.addRoute(mr.router.patch, path, &httpHandler{h: handler}, middleware)
}

// Delete registers a handler for DELETE requests with specific middleware.
func (mr *MiddlewareRouter) Delete(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	middleware := append([]Middleware{}, mr.router.middleware...)
	middleware = append(middleware, mr.middleware...)
	return mr.router.addRoute(mr.router.delete, path, &httpHandler{h: httpHandlerFunc}, middleware)
}

// Delete registers a handler with parameters for DELETE requests with specific middleware.
func (mr *MiddlewareRouter) DeleteParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	middleware := append([]Middleware{}, mr.router.middleware...)
	middleware = append(middleware, mr.middleware...)
	return mr.router.addRoute(mr.router.delete, path, &httpHandler{h: handler}, middleware)
}

// HTTP method registration functions for MiddlewareGroup

// Get registers a handler for GET requests with specific middleware.
func (mg *MiddlewareGroup) Get(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	middleware := append([]Middleware{}, mg.group.router.middleware...)
	middleware = append(middleware, mg.group.middleware...)
	middleware = append(middleware, mg.middleware...)
	return mg.group.router.addRoute(mg.group.router.get, mg.group.prefix+path, &httpHandler{h: httpHandlerFunc}, middleware)
}

// Get registers a handler with parameters for GET requests with specific middleware.
func (mg *MiddlewareGroup) GetParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	middleware := append([]Middleware{}, mg.group.router.middleware...)
	middleware = append(middleware, mg.group.middleware...)
	middleware = append(middleware, mg.middleware...)
	return mg.group.router.addRoute(mg.group.router.get, mg.group.prefix+path, &httpHandler{h: handler}, middleware)
}

// Post registers a handler for POST requests with specific middleware.
func (mg *MiddlewareGroup) Post(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	middleware := append([]Middleware{}, mg.group.router.middleware...)
	middleware = append(middleware, mg.group.middleware...)
	middleware = append(middleware, mg.middleware...)
	return mg.group.router.addRoute(mg.group.router.post, mg.group.prefix+path, &httpHandler{h: httpHandlerFunc}, middleware)
}

// Post registers a handler with parameters for POST requests with specific middleware.
func (mg *MiddlewareGroup) PostParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	middleware := append([]Middleware{}, mg.group.router.middleware...)
	middleware = append(middleware, mg.group.middleware...)
	middleware = append(middleware, mg.middleware...)
	return mg.group.router.addRoute(mg.group.router.post, mg.group.prefix+path, &httpHandler{h: handler}, middleware)
}

// Put registers a handler for PUT requests with specific middleware.
func (mg *MiddlewareGroup) Put(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	middleware := append([]Middleware{}, mg.group.router.middleware...)
	middleware = append(middleware, mg.group.middleware...)
	middleware = append(middleware, mg.middleware...)
	return mg.group.router.addRoute(mg.group.router.put, mg.group.prefix+path, &httpHandler{h: httpHandlerFunc}, middleware)
}

// Put registers a handler with parameters for PUT requests with specific middleware.
func (mg *MiddlewareGroup) PutParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	middleware := append([]Middleware{}, mg.group.router.middleware...)
	middleware = append(middleware, mg.group.middleware...)
	middleware = append(middleware, mg.middleware...)
	return mg.group.router.addRoute(mg.group.router.put, mg.group.prefix+path, &httpHandler{h: handler}, middleware)
}

// Patch registers a handler for PATCH requests with specific middleware.
func (mg *MiddlewareGroup) Patch(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	middleware := append([]Middleware{}, mg.group.router.middleware...)
	middleware = append(middleware, mg.group.middleware...)
	middleware = append(middleware, mg.middleware...)
	return mg.group.router.addRoute(mg.group.router.patch, mg.group.prefix+path, &httpHandler{h: httpHandlerFunc}, middleware)
}

// Patch registers a handler with parameters for PATCH requests with specific middleware.
func (mg *MiddlewareGroup) PatchParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	middleware := append([]Middleware{}, mg.group.router.middleware...)
	middleware = append(middleware, mg.group.middleware...)
	middleware = append(middleware, mg.middleware...)
	return mg.group.router.addRoute(mg.group.router.patch, mg.group.prefix+path, &httpHandler{h: handler}, middleware)
}

// Delete registers a handler for DELETE requests with specific middleware.
func (mg *MiddlewareGroup) Delete(path string, handler http.HandlerFunc) error {
	httpHandlerFunc := func(w http.ResponseWriter, r *http.Request, _ []string) {
		handler(w, r)
	}

	middleware := append([]Middleware{}, mg.group.router.middleware...)
	middleware = append(middleware, mg.group.middleware...)
	middleware = append(middleware, mg.middleware...)
	return mg.group.router.addRoute(mg.group.router.delete, mg.group.prefix+path, &httpHandler{h: httpHandlerFunc}, middleware)
}

// Delete registers a handler with parameters for DELETE requests with specific middleware.
func (mg *MiddlewareGroup) DeleteParam(path string, handler func(w http.ResponseWriter, r *http.Request, params []string)) error {
	middleware := append([]Middleware{}, mg.group.router.middleware...)
	middleware = append(middleware, mg.group.middleware...)
	middleware = append(middleware, mg.middleware...)
	return mg.group.router.addRoute(mg.group.router.delete, mg.group.prefix+path, &httpHandler{h: handler}, middleware)
}

// readSegment extracts the next path segment starting at the given position.
// Returns the segment, the position after it, and whether there are more segments.
func readSegment(path string, start int) (segment string, end int, hasMore bool) {
	if start >= len(path) {
		return "", start, false
	}

	if path[start] == '/' {
		start++
	}

	if start >= len(path) {
		return "", start, false
	}

	end = start
	for end < len(path) && path[end] != '/' {
		end++
	}

	return path[start:end], end, end < len(path)
}

// addRoute adds a new route to the prefix tree with middleware.
func (r *Router) addRoute(root *node, path string, handler Handler, middleware []Middleware) error {
	wrappedHandler := applyMiddleware(handler, middleware)

	if path == "/" {
		root.handler = wrappedHandler
		return nil
	}

	current := root
	pos := 0
	var lastWildcard bool
	paramsCount := uint8(0)

	for {
		segment, newPos, hasMore := readSegment(path, pos)
		if segment == "" {
			break
		}

		isDynamic := len(segment) > 2 && segment[0] == '[' && segment[len(segment)-1] == ']'
		isWildcard := len(segment) > 0 && segment[0] == '*'

		if isWildcard {
			if lastWildcard {
				return fmt.Errorf("wildcard must be the last segment in the path")
			}
			if hasMore {
				return fmt.Errorf("wildcard must be the last segment in the path")
			}
			lastWildcard = true
		}

		if isDynamic || isWildcard {
			paramsCount++
		}

		var child *node
		for _, n := range current.children {
			if n.segment == segment {
				child = n
				break
			}
		}

		if child == nil {
			child = &node{
				segment:    segment,
				isDynamic:  isDynamic,
				isWildcard: isWildcard,
			}
			current.children = append(current.children, child)
		}

		if child.maxParams < paramsCount {
			child.maxParams = paramsCount
		}
		current = child
		pos = newPos
	}

	current.handler = wrappedHandler
	return nil
}

// Lookup finds a handler matching the given method and path.
// Returns the handler, any captured parameters, and whether a match was found.
func (r *Router) Lookup(method, path string) (Handler, []string, bool) {
	var root *node
	switch method {
	case "GET":
		root = r.get
	case "POST":
		root = r.post
	case "PUT":
		root = r.put
	case "PATCH":
		root = r.patch
	case "DELETE":
		root = r.delete
	default:
		return nil, nil, false
	}

	if path == "/" {
		return root.handler, nil, root.handler != nil
	}

	params := make([]string, 0, root.maxParams)
	h, found := match(root, path, 0, &params)
	if !found {
		return nil, nil, false
	}
	return h, params, true
}

// match recursively traverses the prefix tree to find a matching handler.
// It populates params with any captured path parameters or wildcard matches.
func match(current *node, path string, start int, params *[]string) (Handler, bool) {
	// Check for wildcard children first
	for _, child := range current.children {
		if child.isWildcard {
			remaining := path[start:]
			if len(remaining) > 0 && remaining[0] == '/' {
				remaining = remaining[1:]
			}
			*params = append(*params, remaining)
			return child.handler, child.handler != nil
		}
	}

	// Read current segment
	segment, pos, hasMore := readSegment(path, start)
	if segment == "" {
		return current.handler, current.handler != nil
	}

	// Try to match children
	for _, child := range current.children {
		if child.segment == segment || child.isDynamic {
			if child.isDynamic {
				*params = append(*params, segment)
			}
			if !hasMore {
				return child.handler, child.handler != nil
			}
			if h, found := match(child, path, pos, params); found {
				return h, true
			}
		}
	}

	return nil, false
}