// weircon-random-proxy — HTTP fetch service that forwards client requests
// through one of N upstream SOCKS5 endpoints (each typically backing a
// distinct WireGuard tunnel).
//
// Intended to sit behind a reverse proxy (e.g. NginxProxyManager) that
// validates an API key header and strips it before forwarding.
package main

import (
	"context"
	_ "embed"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strconv"
	"strings"
	"time"

	"golang.org/x/net/proxy"
)

//go:embed ui.html
var uiHTML []byte

type Config struct {
	Listen      string
	AdminListen string
	ProxyAddrs  []string
	Timeout     time.Duration
	UIEnabled   bool

	// Rotation drain tuning.
	DrainTimeout time.Duration // max wait for in-flight to clear during drain
	DrainSettle  time.Duration // quiet period after in-flight hits zero (the "before")
	DrainWarmup  time.Duration // post-undrain cool-in before re-selection (the "after")
	DrainMaxHold time.Duration // safety: auto-undrain a slot if undrain never arrives

	// Egress-IP prober (feeds /status).
	ProbeURL      string        // IP-echo endpoint fetched through each tunnel
	ProbeInterval time.Duration // 0 disables probing
	ProbeTimeout  time.Duration // per-probe ceiling
}

func loadConfig() (Config, error) {
	cfg := Config{
		Listen:        envStr("WEIRCON_LISTEN", ":8080"),
		AdminListen:   envStr("WEIRCON_ADMIN_LISTEN", "127.0.0.1:8081"),
		Timeout:       time.Duration(envInt("WEIRCON_REQUEST_TIMEOUT_SEC", 30)) * time.Second,
		UIEnabled:     envBool("WEIRCON_UI_ENABLED", true),
		DrainTimeout:  time.Duration(envInt("WEIRCON_DRAIN_TIMEOUT_SEC", 25)) * time.Second,
		DrainSettle:   time.Duration(envInt("WEIRCON_DRAIN_SETTLE_SEC", 2)) * time.Second,
		DrainWarmup:   time.Duration(envInt("WEIRCON_DRAIN_WARMUP_SEC", 3)) * time.Second,
		DrainMaxHold:  time.Duration(envInt("WEIRCON_DRAIN_MAX_HOLD_SEC", 60)) * time.Second,
		ProbeURL:      envStr("WEIRCON_PROBE_URL", "https://api.ipify.org"),
		ProbeInterval: time.Duration(envInt("WEIRCON_PROBE_INTERVAL_SEC", 60)) * time.Second,
		ProbeTimeout:  time.Duration(envInt("WEIRCON_PROBE_TIMEOUT_SEC", 10)) * time.Second,
	}
	if raw := envStr("WEIRCON_PROXY_ADDRS", ""); raw != "" {
		for _, a := range strings.Split(raw, ",") {
			a = strings.TrimSpace(a)
			if a == "" {
				continue
			}
			cfg.ProxyAddrs = append(cfg.ProxyAddrs, a)
		}
	} else {
		host := envStr("WEIRCON_PROXY_HOST", "127.0.0.1")
		base := envInt("WEIRCON_PROXY_BASE_PORT", 25400)
		count := envInt("WEIRCON_PROXY_COUNT", 10)
		for i := 0; i < count; i++ {
			cfg.ProxyAddrs = append(cfg.ProxyAddrs, fmt.Sprintf("%s:%d", host, base+i))
		}
	}
	if len(cfg.ProxyAddrs) == 0 {
		return cfg, errors.New("no proxy addrs configured")
	}
	for i, a := range cfg.ProxyAddrs {
		if _, _, err := net.SplitHostPort(a); err != nil {
			return cfg, fmt.Errorf("invalid proxy addr [%d] %q: %w", i, a, err)
		}
	}
	return cfg, nil
}

func envStr(k, def string) string {
	if v, ok := os.LookupEnv(k); ok && v != "" {
		return v
	}
	return def
}

func envInt(k string, def int) int {
	if v, ok := os.LookupEnv(k); ok && v != "" {
		if n, err := strconv.Atoi(v); err == nil {
			return n
		}
	}
	return def
}

func envBool(k string, def bool) bool {
	v, ok := os.LookupEnv(k)
	if !ok || v == "" {
		return def
	}
	switch strings.ToLower(v) {
	case "1", "true", "yes", "on":
		return true
	case "0", "false", "no", "off":
		return false
	}
	return def
}

// buildTransports creates one *http.Transport per backend SOCKS5 tunnel so
// each egress maintains its own connection pool.
func buildTransports(cfg Config) ([]*http.Transport, error) {
	pool := make([]*http.Transport, len(cfg.ProxyAddrs))
	for i, addr := range cfg.ProxyAddrs {
		base := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}
		d, err := proxy.SOCKS5("tcp", addr, nil, base)
		if err != nil {
			return nil, fmt.Errorf("socks5 dialer proxy%d: %w", i, err)
		}
		ctxDialer, ok := d.(proxy.ContextDialer)
		if !ok {
			return nil, fmt.Errorf("socks5 dialer proxy%d lacks ContextDialer", i)
		}
		pool[i] = &http.Transport{
			DialContext:           ctxDialer.DialContext,
			MaxIdleConns:          100,
			MaxIdleConnsPerHost:   10,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   15 * time.Second,
			ResponseHeaderTimeout: cfg.Timeout,
			ExpectContinueTimeout: 1 * time.Second,
		}
	}
	return pool, nil
}

type ctxKey struct{}

type reqCtx struct {
	proxyID int
	target  *url.URL
	method  string
}

// poolTransport routes each RoundTrip to the per-tunnel transport carried in
// the request context.
type poolTransport struct{ pool []*http.Transport }

func (p *poolTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	rc, ok := req.Context().Value(ctxKey{}).(reqCtx)
	if !ok {
		return nil, errors.New("internal: missing reqCtx")
	}
	return p.pool[rc.proxyID].RoundTrip(req)
}

// stripHeaders are removed from the outbound request so nothing about the
// caller or the proxy chain can reach the target. It covers our own control
// headers plus every standard forwarding / client-IP header that a fronting
// reverse proxy (NPM, nginx, Cloudflare, …) might inject. Any other
// "X-Forwarded-*" header is swept by prefix in stripIdentifying.
var stripHeaders = []string{
	// weircon control headers
	"X-Weircon-Random-Ip",
	"X-Weircon-Random-Ip-Redirect",
	"X-Weircon-Proxy-Id",
	"X-Weircon-Forward-Method",
	// generic forwarding / origin-IP headers
	"Forwarded",
	"Via",
	"X-Real-Ip",
	"X-Original-Forwarded-For",
	"X-Client-Ip",
	"X-Cluster-Client-Ip",
	"X-Original-Url",
	"X-Original-Host",
	"X-Rewrite-Url",
	"X-Proxy-Id",
	// CDN / vendor client-IP headers
	"Cf-Connecting-Ip",
	"Cf-Ipcountry",
	"Cf-Ray",
	"True-Client-Ip",
	"Fastly-Client-Ip",
	"Fly-Client-Ip",
	"X-Appengine-User-Ip",
}

// stripIdentifying removes every header that could reveal the caller or the
// proxy chain. The explicit list catches the common ones; the prefix sweep
// catches any vendor-specific X-Forwarded-* we didn't enumerate. Header keys
// in an http.Header are already canonicalized, so prefix matching is exact.
func stripIdentifying(h http.Header) {
	for _, k := range stripHeaders {
		h.Del(k)
	}
	for k := range h {
		if strings.HasPrefix(k, "X-Forwarded") {
			h.Del(k)
		}
	}
}

// parsePin reads the optional X-WEIRCON-PROXY-ID header. It returns (-1, nil)
// when no pin is requested (caller should pick at random). The actual slot
// reservation — and the draining check — happens in the manager, not here.
func parsePin(req *http.Request, n int) (int, error) {
	raw := req.Header.Get("X-Weircon-Proxy-Id")
	if raw == "" {
		return -1, nil
	}
	id, err := strconv.Atoi(raw)
	if err != nil {
		return 0, fmt.Errorf("X-WEIRCON-PROXY-ID not an int: %w", err)
	}
	if id < 0 || id >= n {
		return 0, fmt.Errorf("X-WEIRCON-PROXY-ID out of range (0-%d)", n-1)
	}
	return id, nil
}

func parseTarget(req *http.Request) (*url.URL, error) {
	raw := req.Header.Get("X-Weircon-Random-Ip-Redirect")
	if raw == "" {
		raw = req.URL.Query().Get("url")
	}
	if raw == "" {
		return nil, errors.New("missing target — set X-WEIRCON-RANDOM-IP-REDIRECT or ?url=")
	}
	u, err := url.Parse(raw)
	if err != nil {
		return nil, fmt.Errorf("invalid target URL: %w", err)
	}
	if u.Scheme != "http" && u.Scheme != "https" {
		return nil, errors.New("target URL must be http(s)")
	}
	if u.Host == "" {
		return nil, errors.New("target URL missing host")
	}
	return u, nil
}

func resolveMethod(req *http.Request) string {
	if m := req.Header.Get("X-Weircon-Forward-Method"); m != "" {
		return strings.ToUpper(m)
	}
	return req.Method
}

func newHandler(cfg Config, pool []*http.Transport, mgr *proxyManager) http.Handler {
	rp := &httputil.ReverseProxy{
		Rewrite: func(pr *httputil.ProxyRequest) {
			rc, ok := pr.In.Context().Value(ctxKey{}).(reqCtx)
			if !ok {
				return
			}
			pr.Out.URL = rc.target
			pr.Out.Host = rc.target.Host
			pr.Out.Method = rc.method
			pr.Out.RequestURI = ""
			stripIdentifying(pr.Out.Header)
		},
		Transport: &poolTransport{pool: pool},
		ModifyResponse: func(resp *http.Response) error {
			if rc, ok := resp.Request.Context().Value(ctxKey{}).(reqCtx); ok {
				resp.Header.Set("X-Weircon-Egress-Proxy", strconv.Itoa(rc.proxyID))
			}
			return nil
		},
		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
			id := -1
			if rc, ok := r.Context().Value(ctxKey{}).(reqCtx); ok {
				id = rc.proxyID
			}
			log.Printf("upstream error via proxy%d: %v", id, err)
			http.Error(w, fmt.Sprintf("upstream via proxy%d failed: %v", id, err), http.StatusBadGateway)
		},
	}

	count := len(pool)
	mux := http.NewServeMux()
	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintf(w, `{"ok":true,"proxies":%d}`, count)
	})
	// Poll-friendly status for scraper-side dashboards: per-slot egress IP +
	// state, plus when the next rotation is due and which egress it replaces.
	mux.HandleFunc("/status", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.Header().Set("Cache-Control", "no-store")
		_ = json.NewEncoder(w).Encode(mgr.snapshot())
	})
	if cfg.UIEnabled {
		mux.HandleFunc("/ui", func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "text/html; charset=utf-8")
			w.Header().Set("Cache-Control", "no-cache")
			_, _ = w.Write(uiHTML)
		})
	}
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/" {
			http.NotFound(w, r)
			return
		}
		// Friendly hint if someone GETs / in a browser with no target.
		if r.Method == http.MethodGet && r.URL.RawQuery == "" &&
			r.Header.Get("X-Weircon-Random-Ip-Redirect") == "" &&
			strings.Contains(r.Header.Get("Accept"), "text/html") &&
			cfg.UIEnabled {
			http.Redirect(w, r, "/ui", http.StatusFound)
			return
		}
		target, err := parseTarget(r)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
		pin, err := parsePin(r, count)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		// Reserve a slot. A pinned slot that is mid-rotation is rejected (the
		// caller asked for that specific egress); a random pick just skips
		// draining/warming slots. release() runs after the response finishes
		// streaming so drain() can observe a true in-flight count.
		var id int
		if pin >= 0 {
			id = pin
			if err := mgr.acquirePinned(id); err != nil {
				w.Header().Set("Retry-After", "5")
				http.Error(w, err.Error(), http.StatusServiceUnavailable)
				return
			}
		} else {
			id, err = mgr.acquireAny()
			if err != nil {
				w.Header().Set("Retry-After", "5")
				http.Error(w, err.Error(), http.StatusServiceUnavailable)
				return
			}
		}
		defer mgr.release(id)

		ctx := context.WithValue(r.Context(), ctxKey{}, reqCtx{
			proxyID: id,
			target:  target,
			method:  resolveMethod(r),
		})
		rp.ServeHTTP(w, r.WithContext(ctx))
	})
	return mux
}

func main() {
	cfg, err := loadConfig()
	if err != nil {
		log.Fatalf("config: %v", err)
	}
	pool, err := buildTransports(cfg)
	if err != nil {
		log.Fatalf("transports: %v", err)
	}
	mgr := newProxyManager(len(pool), cfg.DrainSettle, cfg.DrainWarmup, cfg.DrainMaxHold)

	// Background egress-IP prober feeds /status. A slot is re-probed shortly
	// after it rotates back in so its new egress IP appears promptly rather
	// than waiting for the next periodic pass.
	if cfg.ProbeInterval > 0 {
		pr := &prober{mgr: mgr, pool: pool, url: cfg.ProbeURL, timeout: cfg.ProbeTimeout}
		mgr.afterUndrain = func(id int) {
			time.AfterFunc(cfg.DrainWarmup+500*time.Millisecond, func() { pr.probe(id) })
		}
		go pr.run(context.Background(), cfg.ProbeInterval)
		log.Printf("egress prober: %s every %s", cfg.ProbeURL, cfg.ProbeInterval)
	}

	// Admin server: localhost-only by default so drain/undrain is reachable by
	// the rotation tooling running in the same LXC but never from the public
	// reverse proxy. Empty WEIRCON_ADMIN_LISTEN disables it entirely.
	if cfg.AdminListen != "" {
		admin := &http.Server{
			Addr:              cfg.AdminListen,
			Handler:           newAdminHandler(cfg, mgr),
			ReadHeaderTimeout: 10 * time.Second,
		}
		go func() {
			log.Printf("admin listening on %s", cfg.AdminListen)
			if err := admin.ListenAndServe(); err != nil {
				log.Fatalf("admin server: %v", err)
			}
		}()
	}

	srv := &http.Server{
		Addr:              cfg.Listen,
		Handler:           newHandler(cfg, pool, mgr),
		ReadHeaderTimeout: 10 * time.Second,
		IdleTimeout:       120 * time.Second,
	}
	log.Printf("weircon-random-proxy listening on %s (proxies=%d: %s, ui=%t)",
		cfg.Listen, len(cfg.ProxyAddrs), strings.Join(cfg.ProxyAddrs, ","), cfg.UIEnabled)
	if err := srv.ListenAndServe(); err != nil {
		log.Fatalf("server: %v", err)
	}
}