package idp

import (
	"context"
	"crypto/subtle"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"errors"
	"log/slog"
	"net/http"
	"net/url"
	"time"
	"unicode/utf8"

	dexstorage "github.com/dexidp/dex/storage"

	"tug/internal/storage"
)

const (
	TokenBodyLimit   = 64 << 10
	reuseRevokeLimit = 5 * time.Second
)

type refreshRevoker interface {
	RevokeRefreshTokenAudited(ctx context.Context, id string, audit storage.AppendAuditInput) (bool, error)
}

type reuseGuard struct {
	next         http.Handler
	dex          dexstorage.Storage
	authenticate func(id, secret string) bool
	revoker      refreshRevoker
	logger       *slog.Logger
}

func revokeOnReuse(next http.Handler, dex dexstorage.Storage, clients []dexstorage.Client, revoker refreshRevoker, logger *slog.Logger) http.Handler {
	return &reuseGuard{next: next, dex: dex, authenticate: clientAuthenticator(clients), revoker: revoker, logger: logger}
}

func clientAuthenticator(clients []dexstorage.Client) func(id, secret string) bool {
	byID := make(map[string]dexstorage.Client, len(clients))
	for _, c := range clients {
		byID[c.ID] = c
	}
	return func(id, secret string) bool {
		c, ok := byID[id]
		if !ok {
			return false
		}
		if c.Public {
			return true
		}
		return c.Secret != "" && subtle.ConstantTimeCompare([]byte(c.Secret), []byte(secret)) == 1
	}
}

func (g *reuseGuard) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost || r.URL.Path != "/token" {
		g.next.ServeHTTP(w, r)
		return
	}
	r.Body = http.MaxBytesReader(w, r.Body, TokenBodyLimit)
	if err := r.ParseForm(); err != nil {
		var tooLarge *http.MaxBytesError
		if errors.As(err, &tooLarge) {
			refuseToken(w, http.StatusRequestEntityTooLarge, "The token request body is too large.")
			return
		}
		refuseToken(w, http.StatusBadRequest, "The token request could not be parsed.")
		return
	}
	rec := &statusRecorder{ResponseWriter: w}
	g.next.ServeHTTP(rec, r)
	if rec.status == http.StatusOK || r.PostForm.Get("grant_type") != "refresh_token" {
		return
	}
	client, ok := g.authenticatedClient(r)
	if !ok {
		return
	}
	id, secret, ok := decodeRefreshToken(r.PostForm.Get("refresh_token"))
	if !ok {
		return
	}
	ctx, cancel := context.WithTimeout(context.WithoutCancel(r.Context()), reuseRevokeLimit)
	defer cancel()
	g.revoke(ctx, id, secret, client)
}

func (g *reuseGuard) revoke(ctx context.Context, id, secret, client string) {
	row, err := g.dex.GetRefresh(ctx, id)
	if errors.Is(err, dexstorage.ErrNotFound) {
		return
	}
	if err != nil {
		g.logger.ErrorContext(ctx, "idp: refresh reuse check: read refresh token", "err", err)
		return
	}
	if row.ClientID != client || subtle.ConstantTimeCompare([]byte(row.Token), []byte(secret)) == 1 {
		return
	}
	generation := "older"
	if row.ObsoleteToken != "" && subtle.ConstantTimeCompare([]byte(row.ObsoleteToken), []byte(secret)) == 1 {
		generation = "previous"
	}
	detail, _ := json.Marshal(map[string]string{"client_id": row.ClientID, "generation": generation})
	revoked, err := g.revoker.RevokeRefreshTokenAudited(ctx, row.ID, storage.AppendAuditInput{Actor: "system", Action: storage.ActionRefreshReuseRevoked, Subject: row.Claims.UserID, Detail: string(detail)})
	if err != nil {
		g.logger.ErrorContext(ctx, "idp: refresh reuse: revoke refresh token", "err", err)
		return
	}
	if !revoked {
		return
	}
	g.logger.WarnContext(ctx, "idp: refresh token reused; its family is revoked", "client_id", row.ClientID, "user_id", row.Claims.UserID, "generation", generation)
}

func (g *reuseGuard) authenticatedClient(r *http.Request) (string, bool) {
	id, secret, basic := r.BasicAuth()
	if basic {
		var err error
		if id, err = url.QueryUnescape(id); err != nil {
			return "", false
		}
		if secret, err = url.QueryUnescape(secret); err != nil {
			return "", false
		}
	} else {
		id, secret = r.PostForm.Get("client_id"), r.PostForm.Get("client_secret")
	}
	return id, g.authenticate(id, secret)
}

func refuseToken(w http.ResponseWriter, status int, description string) {
	body, _ := json.Marshal(struct {
		Error       string `json:"error"`
		Description string `json:"error_description"`
	}{"invalid_request", description})
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	w.Write(body)
}

func decodeRefreshToken(raw string) (id, secret string, ok bool) {
	b, err := base64.RawURLEncoding.DecodeString(raw)
	if err != nil || len(b) == 0 {
		return "", "", false
	}
	for len(b) > 0 {
		key, n := binary.Uvarint(b)
		if n <= 0 || key>>3 == 0 {
			return "", "", false
		}
		b = b[n:]
		field, wire := key>>3, key&7
		var value []byte
		switch wire {
		case 0:
			if _, n = binary.Uvarint(b); n <= 0 {
				return "", "", false
			}
			b = b[n:]
		case 1, 5:
			size := 8
			if wire == 5 {
				size = 4
			}
			if len(b) < size {
				return "", "", false
			}
			b = b[size:]
		case 2:
			l, n := binary.Uvarint(b)
			if n <= 0 || l > uint64(len(b)-n) {
				return "", "", false
			}
			value, b = b[n:n+int(l)], b[n+int(l):]
		default:
			return "", "", false
		}
		if field != 1 && field != 2 {
			continue
		}
		if wire != 2 || !utf8.Valid(value) {
			return "", "", false
		}
		if field == 1 {
			id = string(value)
		} else {
			secret = string(value)
		}
	}
	return id, secret, id != ""
}

type statusRecorder struct {
	http.ResponseWriter
	status int
}

func (s *statusRecorder) WriteHeader(code int) {
	if s.status == 0 {
		s.status = code
	}
	s.ResponseWriter.WriteHeader(code)
}

func (s *statusRecorder) Write(b []byte) (int, error) {
	if s.status == 0 {
		s.status = http.StatusOK
	}
	return s.ResponseWriter.Write(b)
}