package s2s2

import (
	"context"
	"net/http"
	"time"
)

type contextKey string

const authenticatedSubjectContextKey = contextKey("RequestSubject")
const currentRequestKey = contextKey("CurrentRequest")
const startTimeContextKey = contextKey("StartTime")

// SetRequestSubject sets the authorized subject in the request scope.
func SetRequestSubject(ctx context.Context, s AuthorizedSubject) context.Context {
	return context.WithValue(ctx, authenticatedSubjectContextKey, s)
}

// RequestSubject returns the authorized subject in the request scope.
//
// Don't use this if PassthroughIfAuthorizationNotPresented is used. This WILL
// panic if the caller does not provide an access token.
func RequestSubject(ctx context.Context) AuthorizedSubject {
	return ctx.Value(authenticatedSubjectContextKey).(AuthorizedSubject)
}

// OptionalRequestSubject returns the authorized subject in the request scope
// and if the AuthorizedSubject is validated.
//
// This is useful PassthroughIfAuthorizationNotPresented is used.
func OptionalRequestSubject(ctx context.Context) (AuthorizedSubject, bool) {
	authz, ok := ctx.Value(authenticatedSubjectContextKey).(AuthorizedSubject)
	return authz, ok
}

type httpRequest struct {
	AmazonTraceID  string
	RemoteAddress  string
	ForwardedFor   string
	ForwardedPort  string
	ForwardedProto string
	Host           string
}

func setCurrentRequest(ctx context.Context, r *http.Request) context.Context {
	return context.WithValue(ctx, currentRequestKey, &httpRequest{
		AmazonTraceID:  r.Header.Get("X-Amzn-Trace-Id"),
		RemoteAddress:  r.RemoteAddr,
		ForwardedFor:   r.Header.Get("X-Forwarded-For"),
		ForwardedPort:  r.Header.Get("X-Forwarded-Port"),
		ForwardedProto: r.Header.Get("X-Forwarded-Proto"),
		Host:           r.Host,
	})
}

func currentRequest(ctx context.Context) (*httpRequest, bool) {
	if value, ok := ctx.Value(currentRequestKey).(*httpRequest); ok {
		return value, true
	}
	return nil, false
}

// setStartTime sets the time when a request was received.
func setStartTime(ctx context.Context) context.Context {
	return context.WithValue(ctx, startTimeContextKey, time.Now())
}

// startTime returns the time that was set.
func startTime(ctx context.Context) time.Time {
	return ctx.Value(startTimeContextKey).(time.Time)
}
