package mock

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"sync"
	"time"
)

// mockProxyIP is the address every real listener binds to. Instances don't
// get distinct addresses (unlike a real cloud provider): only 127.0.0.1 is
// guaranteed bindable without elevated privileges across platforms — macOS
// does not, by default, route the rest of 127.0.0.0/8 the way Linux does.
// Instances sharing one port are told apart by which listener they share,
// not by IP.
const mockProxyIP = "127.0.0.1"

// connectProxy is a minimal HTTP proxy: it tunnels CONNECT requests
// (hijack + bidirectional copy) and forwards plain absolute-form HTTP
// requests. It exists so the health engine's through-the-proxy probe
// genuinely exercises a CONNECT tunnel against the mock provider, rather
// than the healthcheck being simulated or bypassed for local development.
type connectProxy struct {
	ln net.Listener
	sv *http.Server
}

func newConnectProxy(addr string) (*connectProxy, error) {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, fmt.Errorf("mock proxy: listen %s: %w", addr, err)
	}
	sv := &http.Server{Handler: http.HandlerFunc(handleProxyRequest)}
	go func() {
		// Serve returns http.ErrServerClosed on a clean Close; there is no
		// caller left to report anything else to by the time it returns.
		_ = sv.Serve(ln)
	}()
	return &connectProxy{ln: ln, sv: sv}, nil
}

func (c *connectProxy) close() {
	_ = c.sv.Close()
}

func handleProxyRequest(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodConnect {
		handleConnect(w, r)
		return
	}
	handleForward(w, r)
}

// handleConnect implements the CONNECT tunnel: dial the real destination,
// hijack the client connection, and splice the two together. This is the
// exact mechanism the health engine's default https:// probe URL depends
// on — a proxy that TCP-accepts but can't actually tunnel must fail here,
// not succeed.
func handleConnect(w http.ResponseWriter, r *http.Request) {
	dst, err := net.DialTimeout("tcp", r.Host, 10*time.Second)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	defer dst.Close()

	hijacker, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "hijack unsupported", http.StatusInternalServerError)
		return
	}
	src, buf, err := hijacker.Hijack()
	if err != nil {
		return
	}
	defer src.Close()

	if _, err := src.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n")); err != nil {
		return
	}

	// Any bytes the client already sent past the CONNECT request line
	// before we hijacked are sitting in buf's reader; forward them before
	// starting the raw splice loop.
	if n := buf.Reader.Buffered(); n > 0 {
		if _, err := io.CopyN(dst, buf.Reader, int64(n)); err != nil {
			return
		}
	}

	done := make(chan struct{}, 2)
	go func() { io.Copy(dst, src); done <- struct{}{} }()
	go func() { io.Copy(src, dst); done <- struct{}{} }()
	<-done
}

// handleForward proxies a plain absolute-form HTTP request. CONNECT is the
// path the health engine's default probe exercises, but a probeURL
// override using plain http:// should work too.
func handleForward(w http.ResponseWriter, r *http.Request) {
	outReq := r.Clone(r.Context())
	outReq.RequestURI = ""
	resp, err := http.DefaultTransport.RoundTrip(outReq)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()
	for k, vv := range resp.Header {
		for _, v := range vv {
			w.Header().Add(k, v)
		}
	}
	w.WriteHeader(resp.StatusCode)
	_, _ = io.Copy(w, resp.Body)
}

// sharedProxies tracks one real listener per port, reference-counted
// across every mock.Provider in the process. This is package-level rather
// than a field on Provider because a bound TCP port is a process-global OS
// resource: two independently configured mock-typed provider entries (e.g.
// two named "mock" instances in providers-config.yaml) must not both try
// to bind 127.0.0.1:<port> — the second bind would simply fail. Sharing by
// port, refcounted, means any number of instances across any number of
// Provider values can use the same port safely, and the listener is torn
// down once nothing needs it anymore.
var sharedProxies = struct {
	mu     sync.Mutex
	byPort map[int32]*sharedProxyEntry
}{byPort: make(map[int32]*sharedProxyEntry)}

type sharedProxyEntry struct {
	proxy *connectProxy
	refs  int
}

// acquireProxy returns a release func for a real listener on mockProxyIP:
// port, starting one if this is the first acquire for that port. Safe to
// call concurrently; each returned release func must be called exactly
// once.
func acquireProxy(port int32) (release func(), err error) {
	sharedProxies.mu.Lock()
	defer sharedProxies.mu.Unlock()

	entry, ok := sharedProxies.byPort[port]
	if !ok {
		p, err := newConnectProxy(fmt.Sprintf("%s:%d", mockProxyIP, port))
		if err != nil {
			return nil, err
		}
		entry = &sharedProxyEntry{proxy: p}
		sharedProxies.byPort[port] = entry
	}
	entry.refs++

	var once sync.Once
	release = func() {
		once.Do(func() {
			sharedProxies.mu.Lock()
			defer sharedProxies.mu.Unlock()
			entry.refs--
			if entry.refs <= 0 {
				entry.proxy.close()
				delete(sharedProxies.byPort, port)
			}
		})
	}
	return release, nil
}