package authproxy

import (
	"fmt"
	"html"
	"log"
	"net"
	"net/http"
	"net/url"
	"slices"
	"strings"
	"time"
)

// Server implements the auth endpoints nginx talks to:
//
//	GET /_google-auth/verify    auth_request subrequest: 200 if signed in, 401 otherwise
//	GET /_google-auth/start     begin the OAuth flow (redirects to Google)
//	GET /_google-auth/callback  Google redirect URI (only served on AuthHost)
//	GET /_google-auth/finish    mint the session cookie on the destination app host
//	GET /_google-auth/logout    clear the session cookie
//	GET /_google-auth/healthz   liveness probe
//	GET /_google-auth/          human-readable status page
type Server struct {
	cfg    Config
	box    *box
	nonces *nonceCache
	client *http.Client
}

func New(cfg Config) (*Server, error) {
	b, err := newBox(cfg.CookieSecret)
	if err != nil {
		return nil, err
	}
	return &Server{
		cfg:    cfg,
		box:    b,
		nonces: newNonceCache(),
		client: &http.Client{Timeout: 15 * time.Second},
	}, nil
}

func (s *Server) Routes() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc(RoutePrefix+"/verify", s.handleVerify)
	mux.HandleFunc(RoutePrefix+"/start", s.handleStart)
	mux.HandleFunc(RoutePrefix+"/callback", s.handleCallback)
	mux.HandleFunc(RoutePrefix+"/finish", s.handleFinish)
	mux.HandleFunc(RoutePrefix+"/logout", s.handleLogout)
	mux.HandleFunc(RoutePrefix+"/healthz", s.handleHealthz)
	mux.HandleFunc(RoutePrefix+"/", s.handleStatus)
	mux.HandleFunc("/", s.handleStatus)
	return mux
}

// handleVerify is the nginx auth_request target. Response headers become
// available to nginx as $upstream_http_* variables.
func (s *Server) handleVerify(w http.ResponseWriter, r *http.Request) {
	sess, ok := s.sessionFromRequest(r)
	if !ok {
		w.WriteHeader(http.StatusUnauthorized)
		return
	}
	h := w.Header()
	h.Set("X-Auth-Request-Email", headerSafe(sess.Email))
	h.Set("X-Auth-Request-User", headerSafe(sess.User))
	h.Set("X-Auth-Request-Name", headerSafe(sess.Name))
	h.Set("Cache-Control", "no-store")
	w.WriteHeader(http.StatusOK)
}

// handleStart begins the OAuth flow. nginx proxies unauthenticated requests
// here (via the @google_auth_signin named location) with the original URI in
// the X-Auth-Request-Redirect header.
func (s *Server) handleStart(w http.ResponseWriter, r *http.Request) {
	rd := r.Header.Get("X-Auth-Request-Redirect")
	if rd == "" {
		rd = r.URL.Query().Get("rd")
	}
	rd = sanitizeRedirect(rd)

	// Non-browser clients (API calls, curl) get a clean 401 instead of a
	// redirect chain they can't follow meaningfully.
	if !strings.Contains(r.Header.Get("Accept"), "text/html") {
		w.Header().Set("Content-Type", "application/json")
		w.Header().Set("Cache-Control", "no-store")
		w.WriteHeader(http.StatusUnauthorized)
		fmt.Fprint(w, `{"error":"authentication required"}`)
		return
	}

	host := requestHost(r)
	if host == "" {
		s.htmlError(w, http.StatusBadRequest, "Missing Host header.")
		return
	}

	st := stateClaims{
		Host:  host,
		RD:    rd,
		Proto: s.proto(r),
		Nonce: randToken(),
		Exp:   time.Now().Add(10 * time.Minute).Unix(),
	}
	token, err := s.box.seal("state", st)
	if err != nil {
		s.htmlError(w, http.StatusInternalServerError, "Could not start sign-in.")
		return
	}

	q := url.Values{}
	q.Set("client_id", s.cfg.ClientID)
	q.Set("redirect_uri", s.redirectURI())
	q.Set("response_type", "code")
	q.Set("scope", "openid email profile")
	q.Set("state", token)
	if len(s.cfg.AllowedDomains) == 1 {
		// UX hint only; real enforcement happens in emailAllowed.
		q.Set("hd", s.cfg.AllowedDomains[0])
	}
	http.Redirect(w, r, s.cfg.AuthorizeURL+"?"+q.Encode(), http.StatusFound)
}

// handleCallback is Google's redirect target. It only ever runs on AuthHost,
// exchanges the code, authorizes the email, and bounces the browser back to
// the originating app host with a short-lived hand-off token.
func (s *Server) handleCallback(w http.ResponseWriter, r *http.Request) {
	if requestHost(r) != s.cfg.AuthHost {
		s.htmlError(w, http.StatusNotFound, "This host does not serve the OAuth callback.")
		return
	}
	q := r.URL.Query()
	if errCode := q.Get("error"); errCode != "" {
		s.htmlError(w, http.StatusForbidden, "Google sign-in failed: "+html.EscapeString(errCode))
		return
	}

	var st stateClaims
	if err := s.box.open("state", q.Get("state"), &st); err != nil {
		s.htmlError(w, http.StatusBadRequest, "Invalid sign-in state. Go back to the app and try again.")
		return
	}
	if expired(st.Exp) {
		s.htmlError(w, http.StatusForbidden, "This sign-in attempt expired. Go back to the app and try again.")
		return
	}

	idTok, err := s.exchangeCode(r.Context(), q.Get("code"))
	if err != nil {
		log.Printf("callback: code exchange failed: %v", err)
		s.htmlError(w, http.StatusBadGateway, "Could not complete sign-in with Google. Try again.")
		return
	}
	if err := s.validateIDToken(idTok); err != nil {
		log.Printf("callback: id_token rejected: %v", err)
		s.htmlError(w, http.StatusForbidden, "Google returned an invalid identity token.")
		return
	}
	email := strings.ToLower(idTok.Email)
	if !s.emailAllowed(email) {
		log.Printf("callback: denied %s (not in allowed domains/emails) for host %s", email, st.Host)
		s.htmlError(w, http.StatusForbidden,
			"You are signed in to Google as <b>"+html.EscapeString(email)+"</b>, but that account is not allowed to access this app.")
		return
	}

	hand := handoffClaims{
		Email: email,
		User:  idTok.Sub,
		Name:  idTok.Name,
		Host:  st.Host,
		RD:    st.RD,
		Proto: st.Proto,
		Nonce: randToken(),
		Exp:   time.Now().Add(60 * time.Second).Unix(),
	}
	token, err := s.box.seal("handoff", hand)
	if err != nil {
		s.htmlError(w, http.StatusInternalServerError, "Could not complete sign-in.")
		return
	}
	dest := fmt.Sprintf("%s://%s%s/finish?token=%s", st.Proto, st.Host, RoutePrefix, url.QueryEscape(token))
	http.Redirect(w, r, dest, http.StatusFound)
}

// handleFinish runs on the destination app host and turns a hand-off token
// into a host-scoped session cookie.
func (s *Server) handleFinish(w http.ResponseWriter, r *http.Request) {
	var hand handoffClaims
	if err := s.box.open("handoff", r.URL.Query().Get("token"), &hand); err != nil {
		s.htmlError(w, http.StatusForbidden, "Invalid sign-in token. Go back to the app and try again.")
		return
	}
	if expired(hand.Exp) {
		s.htmlError(w, http.StatusForbidden, "This sign-in token expired. Go back to the app and try again.")
		return
	}
	if hand.Host != requestHost(r) {
		s.htmlError(w, http.StatusForbidden, "This sign-in token was issued for a different host.")
		return
	}
	if !s.nonces.use(hand.Nonce, hand.Exp) {
		s.htmlError(w, http.StatusForbidden, "This sign-in token was already used.")
		return
	}

	sess := sessionClaims{
		Email: hand.Email,
		User:  hand.User,
		Name:  hand.Name,
		Host:  hand.Host,
		Exp:   time.Now().Add(s.cfg.SessionTTL).Unix(),
	}
	value, err := s.box.seal("session", sess)
	if err != nil {
		s.htmlError(w, http.StatusInternalServerError, "Could not create session.")
		return
	}
	http.SetCookie(w, &http.Cookie{
		Name:     s.cfg.CookieName,
		Value:    value,
		Path:     "/",
		MaxAge:   int(s.cfg.SessionTTL.Seconds()),
		HttpOnly: true,
		Secure:   hand.Proto == "https",
		SameSite: http.SameSiteLaxMode,
	})
	log.Printf("signed in %s on %s", sess.Email, sess.Host)
	http.Redirect(w, r, sanitizeRedirect(hand.RD), http.StatusFound)
}

func (s *Server) handleLogout(w http.ResponseWriter, r *http.Request) {
	http.SetCookie(w, &http.Cookie{
		Name:     s.cfg.CookieName,
		Value:    "",
		Path:     "/",
		MaxAge:   -1,
		HttpOnly: true,
		Secure:   s.proto(r) == "https",
		SameSite: http.SameSiteLaxMode,
	})
	if rd := sanitizeRedirect(r.URL.Query().Get("rd")); rd != "/" {
		http.Redirect(w, r, rd, http.StatusFound)
		return
	}
	s.htmlPage(w, http.StatusOK, "Signed out",
		`You have been signed out of <b>`+html.EscapeString(requestHost(r))+`</b>.
		<p><a href="/">Sign in again</a></p>`)
}

func (s *Server) handleHealthz(w http.ResponseWriter, _ *http.Request) {
	w.WriteHeader(http.StatusOK)
	fmt.Fprint(w, "ok")
}

// handleStatus is a small human-readable page for debugging.
func (s *Server) handleStatus(w http.ResponseWriter, r *http.Request) {
	if sess, ok := s.sessionFromRequest(r); ok {
		s.htmlPage(w, http.StatusOK, "Signed in",
			`Signed in to <b>`+html.EscapeString(requestHost(r))+`</b> as <b>`+html.EscapeString(sess.Email)+`</b>`+
				` (`+html.EscapeString(sess.Name)+`).`+
				`<p>Session expires `+time.Unix(sess.Exp, 0).UTC().Format(time.RFC1123)+`.</p>`+
				`<p><a href="`+RoutePrefix+`/logout">Sign out</a></p>`)
		return
	}
	s.htmlPage(w, http.StatusOK, "Not signed in",
		`Not signed in on <b>`+html.EscapeString(requestHost(r))+`</b>.
		<p><a href="`+RoutePrefix+`/start?rd=/">Sign in with Google</a></p>`)
}

// --- helpers ---

func (s *Server) sessionFromRequest(r *http.Request) (*sessionClaims, bool) {
	c, err := r.Cookie(s.cfg.CookieName)
	if err != nil || c.Value == "" {
		return nil, false
	}
	var sess sessionClaims
	if err := s.box.open("session", c.Value, &sess); err != nil {
		return nil, false
	}
	if expired(sess.Exp) {
		return nil, false
	}
	if sess.Host != requestHost(r) {
		return nil, false
	}
	return &sess, true
}

func (s *Server) emailAllowed(email string) bool {
	email = strings.ToLower(email)
	if slices.Contains(s.cfg.AllowedEmails, email) {
		return true
	}
	at := strings.LastIndex(email, "@")
	if at < 0 {
		return false
	}
	return slices.Contains(s.cfg.AllowedDomains, email[at+1:])
}

func (s *Server) redirectURI() string {
	scheme := "https"
	if s.cfg.AllowInsecure {
		scheme = "http"
	}
	return scheme + "://" + s.cfg.AuthHost + RoutePrefix + "/callback"
}

// proto reports the effective client-facing scheme. Unless insecure mode is
// on, everything is treated as https so cookies always carry Secure.
func (s *Server) proto(r *http.Request) string {
	if s.cfg.AllowInsecure && r.Header.Get("X-Forwarded-Proto") == "http" {
		return "http"
	}
	return "https"
}

func requestHost(r *http.Request) string {
	host := r.Host
	if h, _, err := net.SplitHostPort(host); err == nil {
		host = h
	}
	return strings.ToLower(host)
}

// sanitizeRedirect only permits same-host relative paths, preventing open
// redirects. Anything suspicious collapses to "/".
func sanitizeRedirect(rd string) string {
	if rd == "" || !strings.HasPrefix(rd, "/") || strings.HasPrefix(rd, "//") {
		return "/"
	}
	if strings.ContainsAny(rd, "\\\r\n") {
		return "/"
	}
	if strings.HasPrefix(rd, RoutePrefix) {
		return "/"
	}
	return rd
}

func expired(unixSeconds int64) bool {
	return time.Now().Unix() > unixSeconds
}

// headerSafe strips characters that are not safe in an HTTP header value.
func headerSafe(s string) string {
	return strings.Map(func(r rune) rune {
		if r < 32 || r == 127 {
			return -1
		}
		return r
	}, s)
}

func (s *Server) htmlError(w http.ResponseWriter, status int, body string) {
	s.htmlPage(w, status, http.StatusText(status), body)
}

func (s *Server) htmlPage(w http.ResponseWriter, status int, title, body string) {
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.Header().Set("Cache-Control", "no-store")
	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.WriteHeader(status)
	fmt.Fprintf(w, `<!doctype html>
<html><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1">
<title>%s</title>
<style>
body{font-family:system-ui,sans-serif;max-width:36rem;margin:15vh auto 0;padding:0 1rem;color:#222;line-height:1.5}
h1{font-size:1.3rem} a{color:#1a73e8}
</style></head>
<body><h1>%s</h1><p>%s</p></body></html>`, html.EscapeString(title), html.EscapeString(title), body)
}