package tailnet

import (
	"bufio"
	"context"
	"log"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"

	"tailscale.com/derp"
	"tailscale.com/tailcfg"

	"cdr.dev/slog/v3"
	"github.com/coder/websocket"
)

// WithWebsocketSupport returns an http.Handler that upgrades
// connections to the "derp" subprotocol to WebSockets and
// passes them to the DERP server.
// Taken from: https://github.com/tailscale/tailscale/blob/e3211ff88ba85435f70984cf67d9b353f3d650d8/cmd/derper/websocket.go#L21
func WithWebsocketSupport(s *derp.Server, base http.Handler) (http.Handler, func()) {
	var mu sync.Mutex
	var waitGroup sync.WaitGroup
	ctx, cancelFunc := context.WithCancel(context.Background())

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			up := strings.ToLower(r.Header.Get("Upgrade"))

			// Very early versions of Tailscale set "Upgrade: WebSocket" but didn't actually
			// speak WebSockets (they still assumed DERP's binary framing). So to distinguish
			// clients that actually want WebSockets, look for an explicit "derp" subprotocol.
			if up != "websocket" || !strings.Contains(r.Header.Get("Sec-Websocket-Protocol"), "derp") {
				base.ServeHTTP(w, r)
				return
			}

			mu.Lock()
			if ctx.Err() != nil {
				mu.Unlock()
				return
			}
			waitGroup.Add(1)
			mu.Unlock()
			defer waitGroup.Done()
			c, err := websocket.Accept(w, r, &websocket.AcceptOptions{
				Subprotocols:   []string{"derp"},
				OriginPatterns: []string{"*"},
				// Disable compression because we transmit WireGuard messages that
				// are not compressible.
				// Additionally, Safari has a broken implementation of compression
				// (see https://github.com/nhooyr/websocket/issues/218) that makes
				// enabling it actively harmful.
				CompressionMode: websocket.CompressionDisabled,
			})
			if err != nil {
				log.Printf("websocket.Accept: %v", err)
				return
			}
			defer c.Close(websocket.StatusInternalError, "closing")
			if c.Subprotocol() != "derp" {
				c.Close(websocket.StatusPolicyViolation, "client must speak the derp subprotocol")
				return
			}
			wc := websocket.NetConn(ctx, c, websocket.MessageBinary)
			brw := bufio.NewReadWriter(bufio.NewReader(wc), bufio.NewWriter(wc))
			s.Accept(ctx, wc, brw, r.RemoteAddr)
		}), func() {
			cancelFunc()
			mu.Lock()
			waitGroup.Wait()
			mu.Unlock()
		}
}

type DERPMapRewriter interface {
	RewriteDERPMap(derpMap *tailcfg.DERPMap)
}

// RewriteDERPMapDefaultRelay rewrites the DERP map to use the given access URL
// as the "embedded relay" access URL. The passed derp map is modified in place.
//
// This is used by clients and agents to rewrite the default DERP relay to use
// their preferred access URL. Both of these clients can use a different access
// URL than the deployment has configured (with `--access-url`), so we need to
// accommodate that and respect the locally configured access URL.
//
// Note: passed context is only used for logging.
func RewriteDERPMapDefaultRelay(ctx context.Context, logger slog.Logger, derpMap *tailcfg.DERPMap, accessURL *url.URL) {
	if derpMap == nil {
		return
	}

	accessPort := 80
	if accessURL.Scheme == "https" {
		accessPort = 443
	}
	if accessURL.Port() != "" {
		parsedAccessPort, err := strconv.Atoi(accessURL.Port())
		if err != nil {
			// This should never happen because URL.Port() returns the empty string
			// if the port is not valid.
			logger.Critical(ctx, "failed to parse URL port, using default port",
				slog.F("port", accessURL.Port()),
				slog.F("access_url", accessURL))
		} else {
			accessPort = parsedAccessPort
		}
	}

	for _, region := range derpMap.Regions {
		if !region.EmbeddedRelay {
			continue
		}

		for _, node := range region.Nodes {
			if node.STUNOnly {
				continue
			}
			node.HostName = accessURL.Hostname()
			node.DERPPort = accessPort
			node.ForceHTTP = accessURL.Scheme == "http"
		}
	}
}