package handlers

import (
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
)

// NewServiceProxy returns an http.Handler that reverse-proxies requests to
// targetURL, stripping pathPrefix from the request path before forwarding.
// WebSocket upgrade requests are handled via raw TCP hijacking to support
// long-lived connections.
func NewServiceProxy(targetURL, allowedOrigin, pathPrefix string) http.Handler {
	target, err := url.Parse(targetURL)
	if err != nil {
		panic("service proxy: invalid target URL: " + err.Error())
	}

	stripPrefix := func(p string) string {
		s := strings.TrimPrefix(p, pathPrefix)
		if s == "" {
			s = "/"
		}
		return s
	}

	rp := &httputil.ReverseProxy{
		Director: func(req *http.Request) {
			req.URL.Scheme = target.Scheme
			req.URL.Host = target.Host
			req.URL.Path = stripPrefix(req.URL.Path)
			if req.URL.RawPath != "" {
				req.URL.RawPath = stripPrefix(req.URL.RawPath)
			}
		},
		ModifyResponse: func(resp *http.Response) error {
			if swa := resp.Header.Get("Service-Worker-Allowed"); swa != "" {
				resp.Header.Set("Service-Worker-Allowed", swa)
			}
			return nil
		},
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.EqualFold(r.Header.Get("Upgrade"), "websocket") {
			if allowedOrigin != "" && r.Header.Get("Origin") != allowedOrigin {
				http.Error(w, "forbidden", http.StatusForbidden)
				return
			}
			proxyWebSocket(w, r, target, pathPrefix)
			return
		}
		rp.ServeHTTP(w, r)
	})
}

// proxyWebSocket handles WebSocket upgrade via raw TCP hijacking.
func proxyWebSocket(w http.ResponseWriter, r *http.Request, target *url.URL, pathPrefix string) {
	host := target.Host
	if target.Port() == "" {
		switch target.Scheme {
		case "https":
			host += ":443"
		default:
			host += ":80"
		}
	}

	upstream, err := net.Dial("tcp", host)
	if err != nil {
		http.Error(w, "bad gateway", http.StatusBadGateway)
		return
	}
	defer upstream.Close()

	r.URL.Scheme = target.Scheme
	r.URL.Host = target.Host
	stripped := strings.TrimPrefix(r.URL.Path, pathPrefix)
	if stripped == "" {
		stripped = "/"
	}
	r.URL.Path = stripped
	if r.URL.RawPath != "" {
		rawStripped := strings.TrimPrefix(r.URL.RawPath, pathPrefix)
		if rawStripped == "" {
			rawStripped = "/"
		}
		r.URL.RawPath = rawStripped
	}
	r.RequestURI = r.URL.RequestURI()

	if err := r.Write(upstream); err != nil {
		http.Error(w, "bad gateway", http.StatusBadGateway)
		return
	}

	hijacker, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "websocket not supported", http.StatusInternalServerError)
		return
	}
	clientConn, _, err := hijacker.Hijack()
	if err != nil {
		return
	}
	defer clientConn.Close()

	done := make(chan struct{}, 2)
	go func() {
		_, _ = io.Copy(upstream, clientConn)
		done <- struct{}{}
	}()
	go func() {
		_, _ = io.Copy(clientConn, upstream)
		done <- struct{}{}
	}()
	<-done
}