package lb

import (
	"context"
	"crypto/tls"
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	coreapi "k8s.io/api/core/v1"

	"github.com/rajch/weave/common"
)

// Proxy serves the HTTP and HTTPS listeners of the Gateways this node holds. TLS is terminated with the
// listener's certificate, chosen by SNI and swapped in place when its Secret changes (no restart, no rebind).
// Requests are routed by the attached HTTPRoutes in the order the spec prescribes and passed to a ready endpoint
// of the chosen backend, with the client's own address in X-Forwarded-For. Bodies are streamed, never limited.
type Proxy struct {
	mu      sync.Mutex
	servers map[string]*portServer // by "address:port"
	// listen binds an address (swappable in tests)
	listen    func(addr string) (net.Listener, error)
	transport http.RoundTripper
	certs     map[string]*tls.Certificate // parsed certificates, by Secret uid@resourceVersion
	rr        sync.Map                    // backend key -> *uint64, for round robin over its endpoints
}

func NewProxy() *Proxy {
	return &Proxy{
		servers: map[string]*portServer{},
		certs:   map[string]*tls.Certificate{},
		listen:  func(addr string) (net.Listener, error) { return net.Listen("tcp", addr) },
		transport: &http.Transport{
			DialContext:           (&net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}).DialContext,
			MaxIdleConns:          256,
			MaxIdleConnsPerHost:   32,
			IdleConnTimeout:       90 * time.Second,
			ExpectContinueTimeout: time.Second,
		},
	}
}

// portServer is one bound address:port and what it currently serves.
type portServer struct {
	addr  string
	https bool
	ln    net.Listener
	srv   *http.Server
	cfg   atomic.Pointer[portConfig]
}

// portConfig is everything one address:port serves: its listeners, most specific hostname first.
type portConfig struct {
	listeners []*liveListener
}

type liveListener struct {
	spec ServeListener
	cert *tls.Certificate
}

// Apply serves exactly these listeners: ports no longer listed are closed, new ones bound, and every served
// port's routes and certificates replaced at once. It returns, by "namespace/gateway/listener", the listeners
// it could not bind.
func (p *Proxy) Apply(listeners []ServeListener) map[string]error {
	p.mu.Lock()
	defer p.mu.Unlock()
	failed := map[string]error{}
	byAddr := map[string][]ServeListener{}
	for _, l := range listeners {
		a := net.JoinHostPort(l.Address.String(), strconv.Itoa(int(l.Port)))
		byAddr[a] = append(byAddr[a], l)
	}
	for a, s := range p.servers {
		if group, ok := byAddr[a]; !ok || group[0].HTTPS != s.https {
			s.close()
			delete(p.servers, a)
		}
	}
	used := map[string]bool{}
	for a, group := range byAddr {
		cfg := p.compile(group, used)
		s := p.servers[a]
		if s == nil {
			var err error
			if s, err = p.start(a, group[0].HTTPS); err != nil {
				for _, l := range group {
					failed[l.Gateway+"/"+l.Name] = err
				}
				continue
			}
			p.servers[a] = s
		}
		s.cfg.Store(cfg)
	}
	// forget certificates no listener uses any more
	for k := range p.certs {
		if !used[k] {
			delete(p.certs, k)
		}
	}
	return failed
}

// Close stops serving everything.
func (p *Proxy) Close() {
	p.Apply(nil)
}

func (p *Proxy) compile(group []ServeListener, used map[string]bool) *portConfig {
	cfg := &portConfig{}
	for _, l := range group {
		ll := &liveListener{spec: l}
		if l.HTTPS && l.Cert != nil {
			ll.cert = p.cert(l.Cert)
			used[certKey(l.Cert)] = true
		}
		cfg.listeners = append(cfg.listeners, ll)
	}
	// most specific first: exact hostnames, then wildcards (longest first), then the catch-all
	sort.SliceStable(cfg.listeners, func(i, j int) bool {
		return hostRank(cfg.listeners[i].spec.Hostname) > hostRank(cfg.listeners[j].spec.Hostname)
	})
	return cfg
}

func certKey(s *coreapi.Secret) string {
	return string(s.UID) + "@" + s.ResourceVersion + "/" + s.Namespace + "/" + s.Name
}

// cert parses a Secret's certificate once per version of the Secret.
func (p *Proxy) cert(s *coreapi.Secret) *tls.Certificate {
	k := certKey(s)
	if c, ok := p.certs[k]; ok {
		return c
	}
	c, err := tls.X509KeyPair(s.Data[coreapi.TLSCertKey], s.Data[coreapi.TLSPrivateKeyKey])
	if err != nil {
		common.Log.Warnf("lb: gateway certificate %s/%s: %v", s.Namespace, s.Name, err)
		return nil
	}
	p.certs[k] = &c
	return &c
}

// hostRank orders listener hostnames: exact (by length) above wildcards (by length) above none.
func hostRank(h string) int {
	switch {
	case h == "":
		return 0
	case strings.HasPrefix(h, "*."):
		return 1000 + len(h)
	}
	return 100000 + len(h)
}

func (p *Proxy) start(addr string, https bool) (*portServer, error) {
	ln, err := p.listen(addr)
	if err != nil {
		return nil, err
	}
	s := &portServer{addr: addr, https: https, ln: ln}
	s.cfg.Store(&portConfig{})
	s.srv = &http.Server{
		Handler:           http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { p.serve(s, w, r) }),
		ReadHeaderTimeout: 30 * time.Second,
		IdleTimeout:       2 * time.Minute,
		// handshakes from scanners and the like would flood the log
		ErrorLog: log.New(debugWriter{}, "lb: gateway "+addr+": ", 0),
	}
	if https {
		s.srv.TLSConfig = &tls.Config{MinVersion: tls.VersionTLS12, GetCertificate: s.certificate}
		go func() {
			if err := s.srv.ServeTLS(ln, "", ""); err != nil && err != http.ErrServerClosed {
				common.Log.Errorf("lb: gateway %s: %v", addr, err)
			}
		}()
	} else {
		go func() {
			if err := s.srv.Serve(ln); err != nil && err != http.ErrServerClosed {
				common.Log.Errorf("lb: gateway %s: %v", addr, err)
			}
		}()
	}
	common.Log.Infof("lb: gateway serving %s (%s)", addr, map[bool]string{true: "HTTPS", false: "HTTP"}[https])
	return s, nil
}

func (s *portServer) close() {
	srv := s.srv
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		_ = srv.Shutdown(ctx)
		_ = srv.Close()
	}()
	common.Log.Infof("lb: gateway no longer serving %s", s.addr)
}

type debugWriter struct{}

func (debugWriter) Write(b []byte) (int, error) {
	common.Log.Debugf("%s", strings.TrimRight(string(b), "\n"))
	return len(b), nil
}

// listenerFor picks the listener a name (SNI or Host) reaches on a port, most specific hostname first.
func (c *portConfig) listenerFor(name string) *liveListener {
	name = strings.ToLower(name)
	for _, l := range c.listeners {
		if l.spec.Hostname == "" || hostMatches(l.spec.Hostname, name) {
			return l
		}
	}
	return nil
}

// certificate answers a TLS handshake with the certificate of the listener its SNI names.
func (s *portServer) certificate(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
	cfg := s.cfg.Load()
	l := cfg.listenerFor(hello.ServerName)
	if l == nil || l.cert == nil {
		return nil, fmt.Errorf("no listener for %q", hello.ServerName)
	}
	return l.cert, nil
}

// candidate is one match of one rule that a request satisfies, with what decides its precedence.
type candidate struct {
	route   *ServeRoute
	rule    *ServeRule
	ruleIdx int
	// the matched hostname: characters if not a wildcard, and characters
	hostExact, hostLen int
	exact              bool
	prefixLen          int
	method             bool
	headers, query     int
	prefix             string // the path prefix matched (for ReplacePrefixMatch)
}

func (p *Proxy) serve(s *portServer, w http.ResponseWriter, r *http.Request) {
	cfg := s.cfg.Load()
	host := hostOnly(r.Host)
	name := host
	if r.TLS != nil && r.TLS.ServerName != "" {
		name = r.TLS.ServerName
	}
	l := cfg.listenerFor(name)
	if l == nil {
		http.Error(w, "no listener for this host", http.StatusNotFound)
		return
	}
	c := pick(l, host, r)
	if c == nil {
		http.Error(w, "no route", http.StatusNotFound)
		return
	}
	for _, f := range c.rule.Filters {
		if f.Type == filterRequestRedirect && f.RequestRedirect != nil {
			to, code := redirect(r, f.RequestRedirect, l.spec.Port, c.prefix)
			http.Redirect(w, r, to, code)
			return
		}
	}
	b := chooseBackend(c.rule.Backends)
	if b == nil {
		// no resolved backend to send it to: the spec's answer for unresolved backends
		http.Error(w, "no usable backend", http.StatusInternalServerError)
		return
	}
	if len(b.Endpoints) == 0 {
		http.Error(w, "no ready endpoint", http.StatusServiceUnavailable)
		return
	}
	n, _ := p.rr.LoadOrStore(b.Key, new(uint64))
	ep := b.Endpoints[int(atomic.AddUint64(n.(*uint64), 1)-1)%len(b.Endpoints)]
	filters := append(append([]HTTPRouteFilter(nil), c.rule.Filters...), b.Filters...)
	rp := &httputil.ReverseProxy{
		Transport:     p.transport,
		FlushInterval: -1, // stream as it comes (git, server-sent events)
		Rewrite: func(pr *httputil.ProxyRequest) {
			pr.SetURL(&url.URL{Scheme: "http", Host: ep.String()})
			pr.Out.Host = pr.In.Host // forwarded unmodified unless a filter says otherwise
			// SetXForwarded replaces whatever the client sent: only this hop's view is trusted
			pr.SetXForwarded()
			if ip, _, err := net.SplitHostPort(pr.In.RemoteAddr); err == nil {
				pr.Out.Header.Set("X-Real-IP", ip)
			}
			for _, f := range filters {
				switch {
				case f.Type == filterRequestHeaderModifier && f.RequestHeaderModifier != nil:
					modifyHeaders(pr.Out.Header, f.RequestHeaderModifier)
				case f.Type == filterURLRewrite && f.URLRewrite != nil:
					if f.URLRewrite.Hostname != nil {
						pr.Out.Host = *f.URLRewrite.Hostname
					}
					if f.URLRewrite.Path != nil {
						pr.Out.URL.Path = rewritePath(pr.In.URL.Path, f.URLRewrite.Path, c.prefix)
						pr.Out.URL.RawPath = ""
					}
				}
			}
		},
		ModifyResponse: func(resp *http.Response) error {
			for _, f := range filters {
				if f.Type == filterResponseHeaderModifier && f.ResponseHeaderModifier != nil {
					modifyHeaders(resp.Header, f.ResponseHeaderModifier)
				}
			}
			return nil
		},
		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
			common.Log.Debugf("lb: gateway %s -> %s: %v", s.addr, ep, err)
			w.WriteHeader(http.StatusBadGateway)
		},
	}
	rp.ServeHTTP(w, r)
}

// pick is the match that takes the request, by the spec's precedence: the most specific matching hostname, then
// an exact path, the longest prefix, a method, the most header matches, the most query matches, then the oldest
// route, its namespace/name, and the first rule.
func pick(l *liveListener, host string, r *http.Request) *candidate {
	var best *candidate
	for ri := range l.spec.Routes {
		route := &l.spec.Routes[ri]
		hostExact, hostLen, ok := routeHost(route, l.spec.Hostname, host)
		if !ok {
			continue
		}
		for i := range route.Rules {
			rule := &route.Rules[i]
			for _, m := range rule.Matches {
				c, ok := matches(m, r)
				if !ok {
					continue
				}
				c.route, c.rule, c.ruleIdx, c.hostExact, c.hostLen = route, rule, i, hostExact, hostLen
				if best == nil || better(c, *best) {
					cc := c
					best = &cc
				}
			}
		}
	}
	return best
}

func better(a, b candidate) bool {
	switch {
	case a.hostExact != b.hostExact:
		return a.hostExact > b.hostExact
	case a.hostLen != b.hostLen:
		return a.hostLen > b.hostLen
	case a.exact != b.exact:
		return a.exact
	case a.prefixLen != b.prefixLen:
		return a.prefixLen > b.prefixLen
	case a.method != b.method:
		return a.method
	case a.headers != b.headers:
		return a.headers > b.headers
	case a.query != b.query:
		return a.query > b.query
	case !a.route.Created.Equal(b.route.Created):
		return a.route.Created.Before(b.route.Created)
	case a.route.Key != b.route.Key:
		return a.route.Key < b.route.Key
	}
	return a.ruleIdx < b.ruleIdx
}

// routeHost reports whether a route answers this host on its listener, and how specifically.
func routeHost(route *ServeRoute, listenerHost, host string) (int, int, bool) {
	names := route.Hostnames
	if len(names) == 0 {
		if listenerHost == "" {
			return 0, 0, true
		}
		names = []string{listenerHost}
	}
	exact, length, ok := 0, 0, false
	for _, h := range names {
		if !hostMatches(h, host) {
			continue
		}
		ok = true
		e := 0
		if !strings.HasPrefix(h, "*.") {
			e = len(h)
		}
		if e > exact || (e == exact && len(h) > length) {
			exact, length = e, len(h)
		}
	}
	return exact, length, ok
}

func matches(m HTTPRouteMatch, r *http.Request) (candidate, bool) {
	var c candidate
	path, value := r.URL.Path, str(m.Path.Value)
	switch str(m.Path.Type) {
	case pathExact:
		if path != value {
			return c, false
		}
		c.exact = true
		c.prefixLen = len(value)
	default:
		trimmed := strings.TrimSuffix(value, "/")
		if trimmed != "" && path != trimmed && !strings.HasPrefix(path, trimmed+"/") {
			return c, false
		}
		c.prefixLen = len(value)
		c.prefix = value
	}
	if m.Method != nil {
		if r.Method != *m.Method {
			return c, false
		}
		c.method = true
	}
	for _, h := range m.Headers {
		if r.Header.Get(h.Name) != h.Value {
			return c, false
		}
	}
	c.headers = len(m.Headers)
	q := r.URL.Query()
	for _, qp := range m.QueryParams {
		if q.Get(qp.Name) != qp.Value {
			return c, false
		}
	}
	c.query = len(m.QueryParams)
	return c, true
}

// chooseBackend picks a backend by weight (nil when there is none to pick).
func chooseBackend(bs []ServeBackend) *ServeBackend {
	total := int32(0)
	for _, b := range bs {
		if b.Weight > 0 {
			total += b.Weight
		}
	}
	if total == 0 {
		return nil
	}
	n := rand.Int31n(total)
	for i := range bs {
		if bs[i].Weight <= 0 {
			continue
		}
		if n < bs[i].Weight {
			return &bs[i]
		}
		n -= bs[i].Weight
	}
	return nil
}

// redirect is where a RequestRedirect filter sends a request, and with which status.
func redirect(r *http.Request, f *HTTPRequestRedirectFilter, listenerPort int32, prefix string) (string, int) {
	scheme := "http"
	if r.TLS != nil {
		scheme = "https"
	}
	port := int(listenerPort)
	if f.Scheme != nil {
		// a new scheme brings its well-known port, unless a port is given
		scheme = *f.Scheme
		port = map[string]int{"http": 80, "https": 443}[scheme]
	}
	if f.Port != nil {
		port = int(*f.Port)
	}
	host := hostOnly(r.Host)
	if f.Hostname != nil {
		host = *f.Hostname
	}
	if !(scheme == "http" && port == 80) && !(scheme == "https" && port == 443) && port != 0 {
		host = net.JoinHostPort(host, strconv.Itoa(port))
	}
	path := r.URL.Path
	if f.Path != nil {
		path = rewritePath(path, f.Path, prefix)
	}
	code := http.StatusFound
	if f.StatusCode != nil {
		code = *f.StatusCode
	}
	return (&url.URL{Scheme: scheme, Host: host, Path: path, RawQuery: r.URL.RawQuery}).String(), code
}

// rewritePath applies a path modifier: the whole path, or the matched prefix replaced.
func rewritePath(path string, m *HTTPPathModifier, prefix string) string {
	switch m.Type {
	case "ReplaceFullPath":
		return str(m.ReplaceFullPath)
	case "ReplacePrefixMatch":
		rest := strings.TrimPrefix(path, strings.TrimSuffix(prefix, "/"))
		out := strings.TrimSuffix(str(m.ReplacePrefixMatch), "/") + rest
		if !strings.HasPrefix(out, "/") {
			out = "/" + out
		}
		return out
	}
	return path
}

func modifyHeaders(h http.Header, f *HTTPHeaderFilter) {
	for _, s := range f.Set {
		h.Set(s.Name, s.Value)
	}
	for _, a := range f.Add {
		h.Add(a.Name, a.Value)
	}
	for _, name := range f.Remove {
		h.Del(name)
	}
}

// hostOnly is a Host header without its port, lowercased (the spec ignores the port when matching).
func hostOnly(h string) string {
	if host, _, err := net.SplitHostPort(h); err == nil {
		h = host
	}
	return strings.ToLower(strings.TrimSuffix(h, "."))
}