package main

import (
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strings"

	goidentity "github.com/jcmturner/goidentity/v6"
	"github.com/jcmturner/gokrb5/v8/keytab"
	"github.com/jcmturner/gokrb5/v8/spnego"
)

const (
	upstreamURL = "http://gitea-http.gitea.svc.cluster.local:3000"
	listenAddr  = ":4000"
)

// statusCapture wraps http.ResponseWriter to capture the HTTP status code
// written by a downstream handler so callers can inspect it after ServeHTTP.
type statusCapture struct {
	http.ResponseWriter
	status int
}

func (s *statusCapture) WriteHeader(code int) {
	s.status = code
	s.ResponseWriter.WriteHeader(code)
}

func main() {
	keytabPath := os.Getenv("KRB5_KTNAME")
	if keytabPath == "" {
		keytabPath = "/etc/krb5/http.keytab"
	}

	kt, err := keytab.Load(keytabPath)
	if err != nil {
		log.Fatalf("failed to load keytab %s: %v", keytabPath, err)
	}

	upstream, _ := url.Parse(upstreamURL)
	proxy := httputil.NewSingleHostReverseProxy(upstream)
	proxy.Director = func(req *http.Request) {
		req.URL.Scheme = upstream.Scheme
		req.URL.Host = upstream.Host
		if _, ok := req.Header["User-Agent"]; !ok {
			req.Header.Set("User-Agent", "")
		}
	}

	// authedInner runs only when SPNEGO negotiation succeeds.
	// It extracts the client username from the request context and injects
	// X-WEBAUTH-USER for Gitea reverse-proxy auto-login.
	authedInner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// goidentity.AddToHTTPRequestContext is called by SPNEGOKRB5Authenticate on success.
		// FromHTTPRequestContext retrieves the identity set there.
		if id := goidentity.FromHTTPRequestContext(r); id != nil {
			username := strings.SplitN(id.UserName(), "@", 2)[0]
			r.Header.Set("X-WEBAUTH-USER", username)
			log.Printf("SPNEGO auth OK: %s -> %s", id.UserName(), username)
		}
		proxy.ServeHTTP(w, r)
	})

	spnegoHandler := spnego.SPNEGOKRB5Authenticate(authedInner, kt)

	mux := http.NewServeMux()

	mux.HandleFunc("/_healthz", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok"))
	})

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.Header.Get("Authorization"), "Negotiate ") {
			// Client is presenting a Kerberos ticket — run full SPNEGO negotiation.
			// On success: X-WEBAUTH-USER is set by authedInner, Gitea auto-logs in.
			// On failure: gokrb5 returns 401 (invalid ticket — correct to reject).
			rw := &statusCapture{ResponseWriter: w}
			spnegoHandler.ServeHTTP(rw, r)
			if rw.status == http.StatusUnauthorized {
				log.Printf("SPNEGO auth FAILED: from=%s host=%s path=%s token-prefix=%.20s",
					r.RemoteAddr, r.Host, r.URL.Path,
					r.Header.Get("Authorization")[len("Negotiate "):])
			}
			return
		}
		// No Kerberos ticket — pass through so Gitea can show its own login page.
		proxy.ServeHTTP(w, r)
	})

	log.Printf("spnego-proxy listening on %s, upstream %s", listenAddr, upstreamURL)
	log.Fatal(http.ListenAndServe(listenAddr, mux))
}