package browser

import (
	"io"
	"net"
	"net/http"
	"time"

	"webby-builder/internal/security/urlguard"
)

// FilteringProxy is an in-process HTTP/HTTPS forward proxy that validates
// every request URL via urlguard.Validate before forwarding. Chrome is
// wired to send all page-tier network traffic through this proxy via the
// --proxy-server flag, closing the gap where Chrome's own network stack
// (redirects, sub-resources, fetch/XHR, WebSocket) can reach internal
// services that the Go SafeDialer never sees.
//
// One proxy is shared across all browser sessions in a builder process —
// urlguard is stateless and the proxy is just filter + forward.
type FilteringProxy struct {
	listener net.Listener
	server   *http.Server
}

// StartFilteringProxy binds an ephemeral loopback port and serves the
// filtering forward proxy in a background goroutine. Call Close to stop.
func StartFilteringProxy() (*FilteringProxy, error) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, err
	}
	p := &FilteringProxy{listener: l}
	p.server = &http.Server{
		Handler:           http.HandlerFunc(p.handle),
		ReadHeaderTimeout: 10 * time.Second,
		IdleTimeout:       60 * time.Second,
	}
	go func() { _ = p.server.Serve(l) }()
	return p, nil
}

// URL returns the proxy address Chrome should be pointed at via --proxy-server.
// We return the "localhost" hostname (not the 127.0.0.1 literal) so that the
// SSRF defense rule `MAP 127.* ~NOTFOUND` in session.go doesn't blackhole
// Chrome's own filtering proxy. The OS resolver maps localhost → 127.0.0.1
// before any Chrome resolver rule fires, so the connection still goes to the
// same loopback socket. Without this, every webBrowser* call fails with
// ERR_PROXY_CONNECTION_FAILED before the agent even reaches the target site.
func (p *FilteringProxy) URL() string {
	if p == nil || p.listener == nil {
		return ""
	}
	_, port, err := net.SplitHostPort(p.listener.Addr().String())
	if err != nil {
		return "http://" + p.listener.Addr().String()
	}
	return "http://localhost:" + port
}

// Close stops the proxy.
func (p *FilteringProxy) Close() error {
	if p == nil || p.server == nil {
		return nil
	}
	return p.server.Close()
}

func (p *FilteringProxy) handle(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodConnect {
		p.handleConnect(w, r)
		return
	}
	p.handleHTTP(w, r)
}

// handleConnect handles HTTPS via the CONNECT tunnel method. The proxy only
// sees host:port — that's enough for urlguard.Validate to reject private IPs
// before the TLS handshake begins. Once tunneled, the proxy can't inspect
// per-request traffic, but every NEW HTTPS resource the page loads opens a
// fresh CONNECT request that goes through this check again.
func (p *FilteringProxy) handleConnect(w http.ResponseWriter, r *http.Request) {
	host, _, err := net.SplitHostPort(r.Host)
	if err != nil {
		http.Error(w, "bad CONNECT host", http.StatusBadRequest)
		return
	}
	if err := urlguard.Validate("https://" + host); err != nil {
		http.Error(w, "ssrf_blocked: "+err.Error(), http.StatusForbidden)
		return
	}

	dest, err := net.DialTimeout("tcp", r.Host, 10*time.Second)
	if err != nil {
		http.Error(w, "upstream dial failed: "+err.Error(), http.StatusBadGateway)
		return
	}
	hj, ok := w.(http.Hijacker)
	if !ok {
		_ = dest.Close()
		http.Error(w, "hijack not supported", http.StatusInternalServerError)
		return
	}
	src, _, err := hj.Hijack()
	if err != nil {
		_ = dest.Close()
		http.Error(w, "hijack failed: "+err.Error(), http.StatusInternalServerError)
		return
	}
	if _, err := src.Write([]byte("HTTP/1.1 200 OK\r\n\r\n")); err != nil {
		_ = src.Close()
		_ = dest.Close()
		return
	}
	// Bidirectional copy until either side closes.
	go func() {
		_, _ = io.Copy(dest, src)
		_ = dest.Close()
	}()
	_, _ = io.Copy(src, dest)
	_ = src.Close()
}

// handleHTTP handles plain HTTP forward-proxy requests. The full target URL
// is in r.URL; urlguard.Validate gets the complete picture (scheme + host).
// Redirects are validated via the http.Client CheckRedirect hook.
func (p *FilteringProxy) handleHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Scheme == "" || r.URL.Host == "" {
		http.Error(w, "absolute URL required", http.StatusBadRequest)
		return
	}
	if err := urlguard.Validate(r.URL.String()); err != nil {
		http.Error(w, "ssrf_blocked: "+err.Error(), http.StatusForbidden)
		return
	}

	r.RequestURI = ""
	r.Header.Del("Proxy-Connection")

	client := &http.Client{
		Timeout: 30 * time.Second,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			if err := urlguard.Validate(req.URL.String()); err != nil {
				return err
			}
			if len(via) >= 10 {
				return http.ErrUseLastResponse
			}
			return nil
		},
	}
	resp, err := client.Do(r)
	if err != nil {
		http.Error(w, "upstream error: "+err.Error(), http.StatusBadGateway)
		return
	}
	defer func() { _ = 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)
}
