package main

import (
	"bufio"
	"crypto/tls"
	"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:
//   browser → Proxy → MITM TLS → (session token from phone) → endpoint directly → browser
//
// For non-gated hosts:
//   browser → Proxy → internet (transparent, no auth)
type Proxy struct {
	phone   *PhoneClient
	gated   *GatedHosts
	mitm    *MITM
	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.EnsureSession()
		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 {
		p.handleGatedConnect(clientConn, host, target)
	} 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)
}

// handleGatedConnect performs MITM TLS interception for a gated HTTPS host.
// Flow:
//  1. Acquire session token from phone (triggers FIDO2 card interaction).
//  2. Respond 200 to browser.
//  3. Wrap browser connection in TLS using a dynamically issued leaf cert.
//  4. Read each inner HTTP request, inject Authorization: Bearer, forward to
//     the actual server over a fresh TLS connection.
//  5. Write the server response back to the browser.
func (p *Proxy) handleGatedConnect(clientConn net.Conn, hostname, target string) {
	defer clientConn.Close()

	token, err := p.phone.EnsureSession()
	if err != nil {
		fmt.Fprintf(clientConn, "HTTP/1.1 407 Proxy Authentication Required\r\n\r\n")
		log.Printf("gated CONNECT %s: auth failed: %v", hostname, err)
		return
	}

	if _, err := fmt.Fprintf(clientConn, "HTTP/1.1 200 Connection Established\r\n\r\n"); err != nil {
		return
	}

	// Wrap the browser-side connection in TLS (we present a MITM cert).
	tlsClient := tls.Server(clientConn, p.mitm.TLSConfig(hostname))
	if err := tlsClient.Handshake(); err != nil {
		p.logf("MITM TLS handshake for %s: %v", hostname, err)
		return
	}
	defer tlsClient.Close()

	// Serve inner HTTP requests on the now-decrypted browser stream.
	clientBuf := bufio.NewReader(tlsClient)
	for {
		req, err := http.ReadRequest(clientBuf)
		if err != nil {
			if err != io.EOF {
				p.logf("read inner request from browser (%s): %v", hostname, err)
			}
			return
		}

		req.Header.Set("Authorization", "Bearer "+token)
		req.Header.Set("Connection", "close") // no keep-alive to upstream
		stripHopByHop(req.Header)

		if err := p.forwardToUpstream(tlsClient, req, hostname, target); err != nil {
			p.logf("forward to upstream %s: %v", target, err)
			return
		}
	}
}

// forwardToUpstream dials the real server over TLS, sends req, writes the
// response back to clientConn.
func (p *Proxy) forwardToUpstream(clientConn io.Writer, req *http.Request, hostname, target string) error {
	upConn, err := tls.Dial("tcp", target, &tls.Config{
		ServerName: hostname,
		// TODO: pin server CA cert for production; accept any for now.
		InsecureSkipVerify: true,
	})
	if err != nil {
		fmt.Fprintf(clientConn, "HTTP/1.1 502 Bad Gateway\r\n\r\n")
		return fmt.Errorf("dial: %w", err)
	}
	defer upConn.Close()

	// Reconstruct a minimal absolute request for writing on the wire.
	req.URL.Scheme = "https"
	req.URL.Host = hostname
	req.RequestURI = ""

	if err := req.Write(upConn); err != nil {
		return fmt.Errorf("write request: %w", err)
	}

	resp, err := http.ReadResponse(bufio.NewReader(upConn), req)
	if err != nil {
		return fmt.Errorf("read response: %w", err)
	}
	defer resp.Body.Close()

	return resp.Write(clientConn)
}

// 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...)
	}
}