package router

import (
	"strings"

	"github.com/fasthttp/router"
	"github.com/valyala/fasthttp"
)

type HostBucketRouter struct {
	routes        map[string]*router.Router
	bktParam      string
	defaultRouter *router.Router
}

func NewHostBucketRouter(bktParam string) HostBucketRouter {
	return HostBucketRouter{
		routes:   make(map[string]*router.Router),
		bktParam: bktParam,
	}
}

func (hr *HostBucketRouter) Default(router *router.Router) {
	hr.defaultRouter = router
}

func (hr HostBucketRouter) Map(host string, h *router.Router) {
	hr.routes[strings.ToLower(host)] = h
}

func (hr HostBucketRouter) Handler(ctx *fasthttp.RequestCtx) {
	bucket, domain := getBucketDomain(getHost(ctx))
	_ = bucket
	domainRouter, ok := hr.routes[strings.ToLower(domain)]
	if !ok {
		domainRouter = hr.defaultRouter
		if domainRouter == nil {
			ctx.Error(fasthttp.StatusMessage(fasthttp.StatusNotFound), fasthttp.StatusNotFound)
			return
		}
	}

	if bucket != "" {
		ctx.SetUserValue(hr.bktParam, bucket)
	}

	domainRouter.Handler(ctx)
}

func getBucketDomain(host string) (bucket string, domain string) {
	parts := strings.Split(host, ".")
	if len(parts) > 1 {
		return parts[0], strings.Join(parts[1:], ".")
	}
	return "", host
}

// getHost tries its best to return the request host.
// According to section 14.23 of RFC 2616 the Host header
// can include the port number if the default value of 80 is not used.
func getHost(r *fasthttp.RequestCtx) string {
	host := string(r.Host())

	if i := strings.Index(host, ":"); i != -1 {
		host = host[:i]
	}

	return host
}