package main

import (
	"bufio"
	"encoding/base64"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

func getEnvDefault(name, def string) string {
	v := os.Getenv(name)
	if v != "" {
		return v
	}
	return def
}

func parseProxy() *url.URL {
	raw := strings.TrimPrefix(strings.TrimSpace(os.Getenv("PROXY")), "http://")
	if raw == "" {
		log.Fatal("missing PROXY (format: login:password@ip:port)")
	}

	creds, hostPort, ok := strings.Cut(raw, "@")
	if !ok {
		log.Fatal("expected login:password@ip:port")
	}

	login, password, ok := strings.Cut(creds, ":")
	if !ok || strings.TrimSpace(login) == "" || strings.TrimSpace(password) == "" {
		log.Fatal("expected login:password before @")
	}

	hostPort = strings.TrimSpace(hostPort)
	if _, _, err := net.SplitHostPort(hostPort); err != nil {
		log.Fatalf("invalid ip:port: %v", err)
	}

	return &url.URL{Scheme: "http", User: url.UserPassword(login, password), Host: hostPort}
}

func proxyHandler(transport *http.Transport, upstream *url.URL) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/health" {
			w.Write([]byte("ok"))
			return
		}

		if r.Method == http.MethodConnect {
			handleTunneling(w, r, upstream)
			return
		}

		handleHTTP(w, r, transport)
	}
}

func handleHTTP(w http.ResponseWriter, r *http.Request, transport *http.Transport) {
	target := r.URL
	if !target.IsAbs() {
		target = &(*r.URL)
		target.Scheme = "http"
		target.Host = r.Host
	}

	outReq, err := http.NewRequestWithContext(r.Context(), r.Method, target.String(), r.Body)
	if err != nil {
		http.Error(w, "bad request", http.StatusBadRequest)
		return
	}
	outReq.Header = r.Header.Clone()
	outReq.RequestURI = ""

	resp, err := transport.RoundTrip(outReq)
	if err != nil {
		http.Error(w, err.Error(), http.StatusServiceUnavailable)
		return
	}
	defer resp.Body.Close()

	copyHeader(w.Header(), resp.Header)
	w.WriteHeader(resp.StatusCode)
	_, _ = io.Copy(w, resp.Body)
}

func handleTunneling(w http.ResponseWriter, r *http.Request, upstream *url.URL) {
	targetAddr := r.Host
	if _, _, err := net.SplitHostPort(targetAddr); err != nil {
		targetAddr = net.JoinHostPort(targetAddr, "443")
	}

	upstreamConn, err := net.DialTimeout("tcp", upstream.Host, 10*time.Second)
	if err != nil {
		http.Error(w, err.Error(), http.StatusServiceUnavailable)
		return
	}

	if _, err := fmt.Fprintf(
		upstreamConn,
		"CONNECT %s HTTP/1.1\r\nHost: %s\r\nProxy-Authorization: %s\r\n\r\n",
		targetAddr,
		targetAddr,
		basicAuthHeader(upstream.User),
	); err != nil {
		_ = upstreamConn.Close()
		http.Error(w, err.Error(), http.StatusServiceUnavailable)
		return
	}

	br := bufio.NewReader(upstreamConn)
	resp, err := http.ReadResponse(br, &http.Request{Method: http.MethodConnect})
	if err != nil {
		_ = upstreamConn.Close()
		http.Error(w, err.Error(), http.StatusServiceUnavailable)
		return
	}
	if resp.StatusCode != http.StatusOK {
		_ = resp.Body.Close()
		_ = upstreamConn.Close()
		http.Error(w, resp.Status, http.StatusBadGateway)
		return
	}
	_ = resp.Body.Close()

	hijacker, ok := w.(http.Hijacker)
	if !ok {
		_ = upstreamConn.Close()
		http.Error(w, "hijacking not supported", http.StatusInternalServerError)
		return
	}

	clientConn, _, err := hijacker.Hijack()
	if err != nil {
		_ = upstreamConn.Close()
		return
	}

	_, _ = clientConn.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n"))

	go transfer(upstreamConn, clientConn)
	go transfer(clientConn, upstreamConn)
}

func basicAuthHeader(user *url.Userinfo) string {
	if user == nil {
		return ""
	}
	password, _ := user.Password()
	token := base64.StdEncoding.EncodeToString([]byte(user.Username() + ":" + password))
	return "Basic " + token
}

func transfer(dst io.WriteCloser, src io.ReadCloser) {
	defer dst.Close()
	defer src.Close()
	_, _ = io.Copy(dst, src)
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

func main() {
	listenAddr := ":" + getEnvDefault("PORT", "80")
	upstream := parseProxy()

	transport := &http.Transport{Proxy: http.ProxyURL(upstream)}

	log.Printf("Proxy running on %s -> http://%s@%s", listenAddr, upstream.User.Username(), upstream.Host)
	log.Fatal(http.ListenAndServe(listenAddr, proxyHandler(transport, upstream)))
}