// 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 (
	_ "embed"
	"context"
	"errors"
	"fmt"
	"log"
	"math/rand/v2"
	"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
	ProxyAddrs []string
	Timeout    time.Duration
	UIEnabled  bool
}

func loadConfig() (Config, error) {
	cfg := Config{
		Listen:    envStr("WEIRCON_LISTEN", ":8080"),
		Timeout:   time.Duration(envInt("WEIRCON_REQUEST_TIMEOUT_SEC", 30)) * time.Second,
		UIEnabled: envBool("WEIRCON_UI_ENABLED", true),
	}
	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)
}

var weirconHeaders = []string{
	"X-Weircon-Random-Ip",
	"X-Weircon-Random-Ip-Redirect",
	"X-Weircon-Proxy-Id",
	"X-Weircon-Forward-Method",
}

func stripWeircon(h http.Header) {
	for _, k := range weirconHeaders {
		h.Del(k)
	}
}

func pickProxyID(req *http.Request, n int) (int, error) {
	raw := req.Header.Get("X-Weircon-Proxy-Id")
	if raw == "" {
		return rand.IntN(n), 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) 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 = ""
			stripWeircon(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)
	})
	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
		}
		id, err := pickProxyID(r, count)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
		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)
	}
	srv := &http.Server{
		Addr:              cfg.Listen,
		Handler:           newHandler(cfg, pool),
		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)
	}
}