package middleware

import (
	"bufio"
	"context"
	"io"
	"net"
	"net/http"
	"strings"
	"sync"
)

type refreshedTokenResponseContextKey struct{}

type refreshedTokenResponseState struct {
	mu                   sync.RWMutex
	bearerToken          string
	sessionControlCookie *http.Cookie
	suppressed           bool
	committed            bool
	header               http.Header
}

// ProtectRefreshedTokenResponses preserves the original public name while
// delegating to ProtectSessionCredentialResponses.
func ProtectRefreshedTokenResponses(next http.Handler) http.Handler {
	return ProtectSessionCredentialResponses(next)
}

// ProtectSessionCredentialResponses reserves HeaderRefreshedToken and the
// session-control Set-Cookie value at the outer HTTP transport boundary. It
// must wrap the Gin engine and remain outside any response middleware so every
// response commit, including Gin's implicit final commit, passes through the
// guard.
//
// RequireAuth and OptionalAuth only request a refreshed credential when this
// guard is present. Existing hosts that do not install it keep legacy session
// validation behavior and never receive a replacement credential. Routes that
// hijack the connection must register WithoutSessionRenewal before auth because
// bytes written after a successful hijack bypass HTTP response guards.
func ProtectSessionCredentialResponses(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if refreshedTokenResponseGuardPresent(r.Context()) {
			next.ServeHTTP(w, r)
			return
		}

		state := &refreshedTokenResponseState{header: w.Header()}
		guard := &refreshedTokenResponseWriter{
			ResponseWriter: w,
			state:          state,
		}
		guard.protectHeaders()

		ctx := context.WithValue(r.Context(), refreshedTokenResponseContextKey{}, state)
		guardedWriter := http.ResponseWriter(guard)
		if _, ok := w.(http.Pusher); ok {
			guardedWriter = &refreshedTokenPusherResponseWriter{refreshedTokenResponseWriter: guard}
		}
		defer guard.protectHeaders()
		next.ServeHTTP(guardedWriter, r.WithContext(ctx))
	})
}

func refreshedTokenResponseGuardPresent(ctx context.Context) bool {
	return refreshedTokenResponseStateFromContext(ctx) != nil
}

func refreshedTokenResponseCanRenew(ctx context.Context) bool {
	state := refreshedTokenResponseStateFromContext(ctx)
	return state != nil && state.canRenew()
}

func setRefreshedToken(ctx context.Context, token string) {
	if token == "" || !validRefreshedToken(token) {
		return
	}
	if state := refreshedTokenResponseStateFromContext(ctx); state != nil {
		state.setBearerToken(token)
	}
}

func validRefreshedToken(token string) bool {
	return token != "" && token == strings.TrimSpace(token) && !strings.ContainsAny(token, "\r\n")
}

func setSessionControlResponseCookie(ctx context.Context, cookie *http.Cookie) error {
	state := refreshedTokenResponseStateFromContext(ctx)
	if state == nil {
		return errSessionCredentialGuard
	}
	return state.setSessionControlCookie(cookie)
}

func setSessionCredentialResponse(ctx context.Context, bearerToken string, cookie *http.Cookie) error {
	if !validRefreshedToken(bearerToken) || cookie == nil {
		return errInvalidSessionCredentialPair
	}
	state := refreshedTokenResponseStateFromContext(ctx)
	if state == nil {
		return errSessionCredentialGuard
	}
	return state.setSessionCredentials(bearerToken, cookie)
}

func refreshedTokenResponseStateFromContext(ctx context.Context) *refreshedTokenResponseState {
	state, _ := ctx.Value(refreshedTokenResponseContextKey{}).(*refreshedTokenResponseState)
	return state
}

func (s *refreshedTokenResponseState) setBearerToken(token string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.suppressed && !s.committed {
		s.bearerToken = token
		s.protectHeadersLocked()
	}
}

func (s *refreshedTokenResponseState) setSessionControlCookie(cookie *http.Cookie) error {
	if cookie == nil {
		return errInvalidSessionControlProof
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.suppressed || s.committed {
		return errSessionCredentialCommitted
	}
	cloned := *cookie
	s.sessionControlCookie = &cloned
	s.protectHeadersLocked()
	return nil
}

func (s *refreshedTokenResponseState) setSessionCredentials(bearerToken string, cookie *http.Cookie) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.suppressed || s.committed {
		return errSessionCredentialCommitted
	}
	cloned := *cookie
	s.bearerToken = bearerToken
	s.sessionControlCookie = &cloned
	s.protectHeadersLocked()
	return nil
}

func (s *refreshedTokenResponseState) canRenew() bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return !s.suppressed && !s.committed
}

func (s *refreshedTokenResponseState) protectHeaders() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.committed {
		return
	}
	s.protectHeadersLocked()
}

func (s *refreshedTokenResponseState) commitHeaders() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.committed {
		return
	}
	s.protectHeadersLocked()
	s.committed = true
}

func (s *refreshedTokenResponseState) suppress() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.suppressed = true
	s.committed = true
	s.bearerToken = ""
	s.sessionControlCookie = nil
	deleteHeaderEqualFold(s.header, HeaderRefreshedToken)
	deleteSetCookieNamed(s.header, SessionControlCookieName)
}

func (s *refreshedTokenResponseState) suppressInterimCredential(switchingProtocols bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	deleteHeaderEqualFold(s.header, HeaderRefreshedToken)
	deleteSetCookieNamed(s.header, SessionControlCookieName)
	if switchingProtocols {
		s.suppressed = true
		s.committed = true
		s.bearerToken = ""
		s.sessionControlCookie = nil
	}
}

func (s *refreshedTokenResponseState) protectHeadersLocked() {
	deleteHeaderEqualFold(s.header, HeaderRefreshedToken)
	deleteSetCookieNamed(s.header, SessionControlCookieName)
	if s.suppressed {
		return
	}
	hasCredential := false
	if s.bearerToken != "" {
		s.header.Set(HeaderRefreshedToken, s.bearerToken)
		hasCredential = true
	}
	if s.sessionControlCookie != nil {
		s.header.Add("Set-Cookie", s.sessionControlCookie.String())
		hasCredential = true
	}
	if !hasCredential {
		return
	}
	deleteHeaderEqualFold(s.header, headerCacheControl)
	s.header.Set(headerCacheControl, cacheControlNoStore)
}

// refreshedTokenResponseWriter wraps the raw net/http writer below Gin. Gin's
// private response writer delegates both explicit and implicit commits here.
type refreshedTokenResponseWriter struct {
	http.ResponseWriter
	state       *refreshedTokenResponseState
	wroteHeader bool
}

var (
	_ http.ResponseWriter = (*refreshedTokenResponseWriter)(nil)
	_ http.Flusher        = (*refreshedTokenResponseWriter)(nil)
	_ http.Hijacker       = (*refreshedTokenResponseWriter)(nil)
	_ http.CloseNotifier  = (*refreshedTokenResponseWriter)(nil) //nolint:staticcheck // Gin still exposes this compatibility method.
	_ io.ReaderFrom       = (*refreshedTokenResponseWriter)(nil)
	_ io.StringWriter     = (*refreshedTokenResponseWriter)(nil)
	_ http.Pusher         = (*refreshedTokenPusherResponseWriter)(nil)
)

type refreshedTokenPusherResponseWriter struct {
	*refreshedTokenResponseWriter
}

func (w *refreshedTokenPusherResponseWriter) Push(target string, opts *http.PushOptions) error {
	return w.ResponseWriter.(http.Pusher).Push(target, opts)
}

func (w *refreshedTokenResponseWriter) WriteHeader(code int) {
	if code >= 100 && code < 200 {
		w.suppressInterimCredential(code)
		w.ResponseWriter.WriteHeader(code)
		return
	}
	if w.wroteHeader {
		return
	}
	w.state.commitHeaders()
	w.wroteHeader = true
	w.ResponseWriter.WriteHeader(code)
}

func (w *refreshedTokenResponseWriter) Write(data []byte) (int, error) {
	if !w.wroteHeader {
		w.WriteHeader(http.StatusOK)
	}
	return w.ResponseWriter.Write(data)
}

func (w *refreshedTokenResponseWriter) WriteString(data string) (int, error) {
	if !w.wroteHeader {
		w.WriteHeader(http.StatusOK)
	}
	if stringWriter, ok := w.ResponseWriter.(io.StringWriter); ok {
		return stringWriter.WriteString(data)
	}
	return w.ResponseWriter.Write([]byte(data))
}

func (w *refreshedTokenResponseWriter) ReadFrom(src io.Reader) (int64, error) {
	if !w.wroteHeader {
		w.WriteHeader(http.StatusOK)
	}
	if readerFrom, ok := w.ResponseWriter.(io.ReaderFrom); ok {
		return readerFrom.ReadFrom(src)
	}
	return io.Copy(struct{ io.Writer }{w.ResponseWriter}, src)
}

func (w *refreshedTokenResponseWriter) Flush() {
	if !w.wroteHeader {
		w.WriteHeader(http.StatusOK)
	}
	_ = http.NewResponseController(w.ResponseWriter).Flush()
}

func (w *refreshedTokenResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	conn, readWriter, err := http.NewResponseController(w.ResponseWriter).Hijack()
	if err != nil {
		return nil, nil, err
	}
	w.state.suppress()
	w.wroteHeader = true
	return conn, readWriter, nil
}

func (w *refreshedTokenResponseWriter) CloseNotify() <-chan bool {
	if notifier, ok := w.ResponseWriter.(http.CloseNotifier); ok { //nolint:staticcheck // Gin still exposes this compatibility method.
		return notifier.CloseNotify()
	}
	return nil
}

func (w *refreshedTokenResponseWriter) Unwrap() http.ResponseWriter {
	return w.ResponseWriter
}

func (w *refreshedTokenResponseWriter) protectHeaders() {
	w.state.protectHeaders()
}

func (w *refreshedTokenResponseWriter) suppressInterimCredential(code int) {
	switchingProtocols := code == http.StatusSwitchingProtocols
	w.state.suppressInterimCredential(switchingProtocols)
	w.wroteHeader = switchingProtocols
}

func deleteHeaderEqualFold(header http.Header, name string) {
	for key := range header {
		if strings.EqualFold(key, name) {
			delete(header, key)
		}
	}
}

func deleteSetCookieNamed(header http.Header, name string) {
	var values []string
	for key, headerValues := range header {
		if strings.EqualFold(key, "Set-Cookie") {
			values = append(values, headerValues...)
			delete(header, key)
		}
	}
	for _, value := range values {
		pair, _, _ := strings.Cut(value, ";")
		cookieName, _, found := strings.Cut(pair, "=")
		if found && strings.TrimSpace(cookieName) == name {
			continue
		}
		header.Add("Set-Cookie", value)
	}
}