package netutil

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/netip"
	"sync"
	"time"

	"github.com/bluesky-social/indigo/util/ssrf"
)

var relayDrainGrace = 30 * time.Second

type IPResolver interface {
	LookupIP(ctx context.Context, host string) ([]netip.Addr, error)
}

type defaultResolver struct{}

func (defaultResolver) LookupIP(ctx context.Context, host string) ([]netip.Addr, error) {
	ips, err := net.DefaultResolver.LookupIP(ctx, "ip", host)
	if err != nil {
		return nil, err
	}
	addrs := make([]netip.Addr, 0, len(ips))
	for _, ip := range ips {
		if addr, ok := netip.AddrFromSlice(ip); ok {
			addrs = append(addrs, addr.Unmap())
		}
	}
	return addrs, nil
}

type OutboundDialer interface {
	DialContext(ctx context.Context, network, address string) (net.Conn, error)
}

type defaultDialer struct {
	dialer net.Dialer
}

func (d *defaultDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
	return d.dialer.DialContext(ctx, network, address)
}

type SafeConnectProxy struct {
	resolver     IPResolver
	dialer       OutboundDialer
	allowPrivate bool
	server       *http.Server
	listener     net.Listener
	addr         string
	mu           sync.Mutex
	closed       bool
	active       map[net.Conn]struct{}
}

// allowPrivate lets it reach private and loopback addresses, which only a dev stack
// whose names all resolve to loopback wants
func NewSafeConnectProxy(resolver IPResolver, dialer OutboundDialer, allowPrivate bool) *SafeConnectProxy {
	if resolver == nil {
		resolver = defaultResolver{}
	}
	if dialer == nil {
		dialer = SSRFDialer(allowPrivate)
	}
	p := &SafeConnectProxy{
		resolver:     resolver,
		dialer:       dialer,
		allowPrivate: allowPrivate,
		active:       make(map[net.Conn]struct{}),
	}
	return p
}

func (p *SafeConnectProxy) Start() (string, error) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return "", fmt.Errorf("starting safe connect proxy listener: %w", err)
	}
	p.listener = l
	p.addr = l.Addr().String()

	p.server = &http.Server{
		Handler: http.HandlerFunc(p.handleHTTP),
	}

	go func() {
		_ = p.server.Serve(l)
	}()

	return p.addr, nil
}

func (p *SafeConnectProxy) Addr() string {
	return p.addr
}

func (p *SafeConnectProxy) Close() error {
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		return nil
	}
	p.closed = true
	// http.Server drops hijacked conns, so close them explicitly on shutdown
	for conn := range p.active {
		_ = conn.Close()
	}
	server := p.server
	p.mu.Unlock()

	if server != nil {
		return server.Close()
	}
	return nil
}

func (p *SafeConnectProxy) track(conns ...net.Conn) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.closed {
		for _, conn := range conns {
			_ = conn.Close()
		}
		return
	}
	for _, conn := range conns {
		p.active[conn] = struct{}{}
	}
}

func (p *SafeConnectProxy) untrack(conns ...net.Conn) {
	p.mu.Lock()
	defer p.mu.Unlock()
	for _, conn := range conns {
		delete(p.active, conn)
	}
}

func (p *SafeConnectProxy) handleHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodConnect {
		http.Error(w, "Method Not Allowed: only CONNECT method is permitted", http.StatusMethodNotAllowed)
		return
	}

	host, port, err := net.SplitHostPort(r.Host)
	if err != nil {
		host = r.Host
		port = "443"
	}

	if port != "443" {
		http.Error(w, "Forbidden: only port 443 is permitted", http.StatusForbidden)
		return
	}

	ctx, cancel := context.WithTimeout(r.Context(), 10*time.Second)
	defer cancel()

	addrs, err := p.resolveAndVerify(ctx, host)
	if err != nil {
		http.Error(w, fmt.Sprintf("Forbidden: SSRF safety check failed: %v", err), http.StatusForbidden)
		return
	}

	targetConn, err := p.dialAny(ctx, addrs, port)
	if err != nil {
		http.Error(w, fmt.Sprintf("Bad Gateway: dial failed: %v", err), http.StatusBadGateway)
		return
	}
	defer targetConn.Close()

	hj, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "Internal Server Error: hijacking not supported", http.StatusInternalServerError)
		return
	}

	clientConn, _, err := hj.Hijack()
	if err != nil {
		http.Error(w, fmt.Sprintf("Internal Server Error: hijacking failed: %v", err), http.StatusInternalServerError)
		return
	}
	defer clientConn.Close()

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

	p.relay(clientConn, targetConn)
}

func (p *SafeConnectProxy) dialAny(ctx context.Context, addrs []netip.Addr, port string) (net.Conn, error) {
	var errs []error
	for _, addr := range addrs {
		conn, err := p.dialer.DialContext(ctx, "tcp", net.JoinHostPort(addr.String(), port))
		if err == nil {
			return conn, nil
		}
		errs = append(errs, fmt.Errorf("%s: %w", addr, err))
	}
	return nil, errors.Join(errs...)
}

func (p *SafeConnectProxy) relay(clientConn, targetConn net.Conn) {
	p.track(clientConn, targetConn)
	defer p.untrack(clientConn, targetConn)
	defer targetConn.Close()
	defer clientConn.Close()

	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		_, _ = io.Copy(targetConn, clientConn)
		halfClose(targetConn)
	}()
	go func() {
		defer wg.Done()
		_, _ = io.Copy(clientConn, targetConn)
		halfClose(clientConn)
	}()

	drained := make(chan struct{})
	go func() {
		wg.Wait()
		close(drained)
	}()

	select {
	case <-drained:
	case <-time.After(relayDrainGrace):
	}
}

func halfClose(conn net.Conn) {
	if cw, ok := conn.(interface{ CloseWrite() error }); ok {
		_ = cw.CloseWrite()
	}
}

func (p *SafeConnectProxy) resolveAndVerify(ctx context.Context, host string) ([]netip.Addr, error) {
	if addr, err := netip.ParseAddr(host); err == nil {
		addr = addr.Unmap()
		if !p.safe(addr) {
			return nil, fmt.Errorf("direct IP %s is not a safe public address", addr)
		}
		return []netip.Addr{addr}, nil
	}

	resolved, err := p.resolver.LookupIP(ctx, host)
	if err != nil {
		return nil, fmt.Errorf("lookup host %q: %w", host, err)
	}
	if len(resolved) == 0 {
		return nil, fmt.Errorf("host %q resolved to no IP addresses", host)
	}

	addrs := make([]netip.Addr, 0, len(resolved))
	for _, addr := range resolved {
		addr = addr.Unmap()
		if !p.safe(addr) {
			return nil, fmt.Errorf("host %q resolved to unsafe IP %s (all answers must be safe)", host, addr)
		}
		addrs = append(addrs, addr)
	}
	return addrs, nil
}

func (p *SafeConnectProxy) safe(addr netip.Addr) bool {
	return p.allowPrivate || ssrf.IsPublicIPAddress(net.IP(addr.AsSlice()))
}