package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"strings"
	"time"
)

// hopByHop headers that must not be forwarded by a proxy (RFC 7230 §6.1).
var hopByHopHeaders = map[string]bool{
	"connection":          true,
	"keep-alive":          true,
	"proxy-authenticate":  true,
	"proxy-authorization": true,
	"te":                  true,
	"trailers":            true,
	"transfer-encoding":   true,
	"upgrade":             true,
	"proxy-connection":    true, // non-standard but common
}

// Proxy is the HTTP/HTTPS proxy handler.
//
// For plain HTTP requests to gated hosts:
//   browser → Proxy → (session token from phone) → endpoint directly → browser
//
// For HTTPS CONNECT to gated hosts:
//   407 — HTTPS interception is not supported; use plain HTTP for gated endpoints.
//
// For non-gated hosts:
//   browser → Proxy → internet (transparent, no auth)
type Proxy struct {
	phone   *PhoneClient
	gated   *GatedHosts
	verbose bool
}

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodConnect {
		p.handleConnect(w, r)
	} else {
		p.handleHTTP(w, r)
	}
}

// handleHTTP handles plain HTTP proxy requests.
// For gated hosts: acquires a session token from the phone, adds it as
// Authorization: Bearer, then calls the endpoint directly.
func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Host == "" {
		http.Error(w, "not a proxy request", http.StatusBadRequest)
		return
	}

	host := r.URL.Hostname()
	port := r.URL.Port()
	if port == "" {
		port = "80"
	}
	isGated := p.gated.IsGated(host, port)
	p.logf("HTTP %s %s (gated=%v)", r.Method, r.URL, isGated)

	// Build outgoing request. RequestURI must be empty for http.Client/RoundTrip.
	out := r.Clone(r.Context())
	out.RequestURI = ""
	stripHopByHop(out.Header)

	if isGated {
		token, err := p.phone.GetTokenForRequest(r.URL.String(), r.Method)
		if err != nil {
			http.Error(w, "auth: "+err.Error(), http.StatusUnauthorized)
			return
		}
		out.Header.Set("Authorization", "Bearer "+token)
	}

	resp, err := http.DefaultTransport.RoundTrip(out)
	if err != nil {
		http.Error(w, "upstream: "+err.Error(), http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

	copyHeaders(w.Header(), resp.Header)
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}

// handleConnect handles HTTPS CONNECT tunnels.
// For gated hosts: does TLS MITM so Authorization can be injected into each
// inner HTTP request before it is forwarded to the actual server.
// For non-gated hosts: transparent byte-level tunnel.
func (p *Proxy) handleConnect(w http.ResponseWriter, r *http.Request) {
	host, portStr, err := net.SplitHostPort(r.Host)
	if err != nil {
		// No port — default to 443.
		host = r.Host
		portStr = "443"
	}
	if portStr == "" {
		portStr = "443"
	}
	target := net.JoinHostPort(host, portStr)
	isGated := p.gated.IsGated(host, portStr)
	p.logf("CONNECT %s (gated=%v)", target, isGated)

	hijacker, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "hijack not supported", http.StatusInternalServerError)
		return
	}
	clientConn, _, err := hijacker.Hijack()
	if err != nil {
		log.Printf("hijack: %v", err)
		return
	}

	if isGated {
		fmt.Fprintf(clientConn, "HTTP/1.1 407 Proxy Authentication Required\r\nContent-Type: text/plain\r\nProxy-Authenticate: Bearer realm=\"chromecard\"\r\n\r\nHTTPS tunnels to gated hosts are not supported. Use plain HTTP.\r\n")
		p.logf("CONNECT %s: gated HTTPS not supported, returned 407", target)
		clientConn.Close()
	} else {
		p.handleDirectConnect(clientConn, target)
	}
}

// handleDirectConnect tunnels bytes transparently — no auth, no inspection.
func (p *Proxy) handleDirectConnect(clientConn net.Conn, target string) {
	defer clientConn.Close()

	upConn, err := net.DialTimeout("tcp", target, 10*time.Second)
	if err != nil {
		fmt.Fprintf(clientConn, "HTTP/1.1 502 Bad Gateway\r\n\r\n")
		return
	}
	defer upConn.Close()

	fmt.Fprintf(clientConn, "HTTP/1.1 200 Connection Established\r\n\r\n")
	pipe(clientConn, upConn)
}

// pipe copies bytes bidirectionally between two connections until either closes.
func pipe(a, b net.Conn) {
	done := make(chan struct{}, 2)
	go func() { io.Copy(a, b); done <- struct{}{} }()
	go func() { io.Copy(b, a); done <- struct{}{} }()
	<-done
}

// stripHopByHop removes hop-by-hop headers and any headers named in Connection.
func stripHopByHop(h http.Header) {
	if conn := h.Get("Connection"); conn != "" {
		for _, name := range strings.Split(conn, ",") {
			h.Del(strings.TrimSpace(name))
		}
	}
	for name := range hopByHopHeaders {
		h.Del(name)
	}
}

// copyHeaders copies non-hop-by-hop headers from src to dst.
func copyHeaders(dst, src http.Header) {
	for k, vs := range src {
		if !hopByHopHeaders[strings.ToLower(k)] {
			for _, v := range vs {
				dst.Add(k, v)
			}
		}
	}
}

func (p *Proxy) logf(format string, args ...any) {
	if p.verbose {
		log.Printf(format, args...)
	}
}