package handler

import (
	"net/http"
	"net/url"
	"strings"

	"github.com/samber/lo"
)

type ProxyRequestUpdater func(arg ...string) func(r *http.Request)

var ProxyRequestUpdateMap = map[string]ProxyRequestUpdater{
	"HostSchemas": func(arg ...string) func(r *http.Request) {
		targetUrl := arg[0]
		return func(r *http.Request) {
			turl, _ := url.Parse(targetUrl)
			r.URL.Host = turl.Host
			r.URL.Scheme = turl.Scheme
		}
	},
	"HeaderOrigin": func(arg ...string) func(r *http.Request) {

		return func(r *http.Request) {
			r.Header.Set("Origin", r.URL.Scheme+"://"+r.URL.Host)
		}
	},
	"Path": func(arg ...string) func(r *http.Request) {
		replace := arg[0]
		with := arg[1]
		return func(r *http.Request) {
			r.URL.Path = r.URL.Path[len(replace):]
			r.URL.Path = with + r.URL.Path
		}
	},
	"RemoveCookie": func(arg ...string) func(r *http.Request) {
		cookie := arg[0]
		return func(r *http.Request) {
			cookies := r.Cookies()
			_, index, ok := lo.FindIndexOf(cookies, func(v *http.Cookie) bool {
				return v.Name == cookie
			})
			if ok {
				r.Header.Del("Cookie")
				if len(cookies) == 1 {
					return
				}
				cookies = append(cookies[:index], cookies[index+1:]...)
				for _, cookie := range cookies {
					r.AddCookie(cookie)
				}
			}

		}
	},
}

func GetUpdaterFn(directive string) func(r *http.Request) {
	strs := strings.Split(directive, " ")
	if len(strs) > 1 {
		updater, ok := ProxyRequestUpdateMap[strs[0]]
		if ok {
			return updater(strs[1:]...)
		}
	} else {
		updater, ok := ProxyRequestUpdateMap[directive]
		if ok {
			return updater()
		}
	}
	return nil
}