package api

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"slices"
	"strings"
	"time"

	"github.com/lestrrat-go/jwx/v2/jwt"
	"google.golang.org/protobuf/types/known/structpb"
	"google.golang.org/protobuf/types/known/timestamppb"
	"google.golang.org/protobuf/types/known/wrapperspb"

	auditV1 "dev.azure.com/schwarzit/schwarzit.stackit-public/audit-go.git/gen/go/audit/v1"
	pkgAuditCommon "dev.azure.com/schwarzit/schwarzit.stackit-public/audit-go.git/pkg/audit/common"
)

const EmailAddressDoNotReplyAtStackItDotCloud = "do-not-reply@stackit.cloud"
const TokenClaimStackitProjectId = "stackit/project/project.id"
const TokenClaimStackitServiceAccountId = "stackit/serviceaccount/service-account.uid"

var ErrInvalidRequestBody = errors.New("invalid request body")
var ErrInvalidResponse = errors.New("invalid response")
var ErrInvalidAuthorizationHeaderValue = errors.New("invalid authorization header value")
var ErrInvalidBearerToken = errors.New("invalid bearer token")
var ErrTokenIsNotBearerToken = errors.New("token is not a bearer token")

// AuditRequest bundles request related parameters
type AuditRequest struct {

	// The operation request. This may not include all request parameters,
	// such as those that are too large, privacy-sensitive, or duplicated
	// elsewhere in the log record.
	// It should never include user-generated data, such as file contents.
	//
	// Required: true
	Request *pkgAuditCommon.ApiRequest

	// The IP address of the caller.
	// For caller from internet, this will be public IPv4 or IPv6 address.
	// For caller from a VM / K8s Service / etc., this will be the SIT proxy's IPv4 address.
	//
	// Required: true
	RequestClientIP string

	// Correlate multiple audit logs by setting the same id
	//
	// Required: false
	RequestCorrelationId *string

	// The unique ID for a request, which can be propagated to downstream
	// systems. The ID should have low probability of collision
	// within a single day for a specific service.
	//
	// More information can be found here: https://google.aip.dev/155
	//
	// Format: <idempotency-key>
	// Where:
	//   Idempotency-key: Typically consists of an id + version
	//
	// Examples:
	//   5e3952a9-b628-4be6-ac61-b1c6eb4a110c/5
	//
	// Required: false
	RequestId *string

	// The timestamp when the `destination` service receives the first byte of
	// the request.
	//
	// Required: false
	RequestTime *time.Time
}

// AuditResponse bundles response related parameters
type AuditResponse struct {

	// The operation response. This may not include all response elements,
	// such as those that are too large, privacy-sensitive, or duplicated
	// elsewhere in the log record.
	//
	// Required: false
	ResponseBodyBytes []byte

	// The http or gRPC status code.
	//
	// Examples:
	//   https://developer.mozilla.org/en-US/docs/Web/HTTP/Status
	//   https://grpc.github.io/grpc/core/md_doc_statuscodes.html
	//
	// Required: true
	ResponseStatusCode int

	// The HTTP response headers.
	//
	// Required: true
	ResponseHeaders map[string][]string

	// The number of items returned from a List or Query API method,
	// if applicable.
	//
	// Required: false
	ResponseNumItems *int64

	// The timestamp when the "destination" service generates the first byte of
	// the response.
	//
	// Required: false
	ResponseTime *time.Time
}

// AuditMetadata bundles audit event related metadata
type AuditMetadata struct {

	// A unique identifier for the log entry.
	// Is used to check completeness of audit events over time.
	//
	// Format: <unix-timestamp>/<region-zone>/<worker-id>/<sequence-number>
	// Where:
	//   Unix-Timestamp: A UTC unix timestamp in seconds is expected
	//   Region-Zone: The region and (optional) zone id. If both, separated with a - (dash)
	//   Worker-Id: The ID of the K8s Pod, Service-Instance, etc. (must be unique for a sending service)
	//   Sequence-Number: Increasing number, representing the message offset per Worker-Id
	//                    If the Worker-Id changes, the sequence-number has to be reset to 0.
	//
	// Examples:
	// 	 "1721899117/eu01/319a7fb9-edd2-46c6-953a-a724bb377c61/8792726390909855142"
	//
	// Required: true
	AuditInsertId string

	// A set of user-defined (key, value) data that provides additional
	// information about the log entry.
	//
	// Required: false
	AuditLabels *map[string]string

	// The resource name of the log to which this log entry belongs.
	//
	// Format: <pluralType>/<identifier>/logs/<eventType>
	// Where:
	//   Plural-Types: One from the list of supported ObjectType as plural
	//   Event-Types: admin-activity, system-event, policy-denied, data-access
	//
	// Examples:
	//   "projects/00b0f972-59ff-48f2-a4f9-29c57b75c2fa/logs/admin-activity"
	//
	// Required: true
	AuditLogName string

	// The severity of the log entry.
	//
	// Required: true
	AuditLogSeverity auditV1.LogSeverity

	// The name of the service method or operation.
	//
	// Format: stackit.<product>.<version>.<type-chain>.<operation>
	// Where:
	//   Product: The name of the service in lowercase
	//   Version: Optional API version
	//   Type-Chain: Optional chained path to object
	//   Operation: The name of the operation in lowercase
	//
	// Examples:
	//   "stackit.resource-manager.v1.organizations.create"
	//   "stackit.authorization.v1.projects.volumes.create"
	//   "stackit.authorization.v2alpha.projects.volumes.create"
	//   "stackit.authorization.v2.folders.move"
	//   "stackit.resource-manager.health"
	//
	// Required: true
	AuditOperationName string

	// The required IAM permission.
	//
	// Examples:
	//   "resourcemanager.project.edit"
	//
	// Required: false
	AuditPermission *string

	// Result of the IAM permission check.
	//
	// Required: false
	AuditPermissionGranted *bool

	// The resource or collection that is the target of the operation.
	// The name is a scheme-less URI, not including the API service name.
	//
	// Format: <pluralType>/<id>[/locations/<region-zone>][/<details>]
	// Where:
	//   Plural-Type: One from the list of supported ObjectType as plural
	//   Id: The identifier of the object
	//   Region-Zone: Optional region and zone id. If both, separated with a - (dash). Alternatively _ (underscore).
	//   Details: Optional "<key>/<id>" pairs
	//
	// Examples:
	//   "organizations/40ab14ad-b7b0-4b1c-be41-5bc820a968d1"
	//   "projects/7046e7b6-5ae9-441c-99fe-2cd28a5078ec/locations/_/instances/instance-20240723-174217"
	//   "projects/7046e7b6-5ae9-441c-99fe-2cd28a5078ec/locations/eu01/instances/instance-20240723-174217"
	//   "projects/7046e7b6-5ae9-441c-99fe-2cd28a5078ec/locations/sx-stoi01/instances/instance-20240723-174217"
	//   "projects/dd7d1807-54e9-4426-8994-721758b5b554/locations/eu01-m/vms/b6851b4e-7a9d-4973-ab0f-a80a13ee3060/ports/78f8bad4-a291-4fa3-b07f-4a1985d3dbe8"
	//
	// Required: true
	AuditResourceName string

	// The name of the API service performing the operation.
	//
	// Examples:
	//   "resource-manager"
	//
	// Required: true
	AuditServiceName string

	// The time the event described by the log entry occurred.
	//
	// Required: false
	AuditTime *time.Time
}

// NewAuditLogEntry constructs a new audit log event for the given parameters
func NewAuditLogEntry(

	// Required request parameters
	auditRequest AuditRequest,

	// Required response parameters
	auditResponse AuditResponse,

	// Optional map that is added as "details" to the message
	eventMetadata map[string]interface{},

	// Required metadata
	auditMetadata AuditMetadata,
) (*auditV1.AuditLogEntry, error) {

	// Get request headers
	filteredRequestHeaders := FilterAndMergeHeaders(auditRequest.Request.Header)
	filteredResponseHeaders := FilterAndMergeHeaders(auditResponse.ResponseHeaders)

	// Get response body
	responseBody, err := NewResponseBody(auditResponse.ResponseBodyBytes)
	if err != nil {
		return nil, errors.Join(err, ErrInvalidResponse)
	}
	var responseLength *int64
	if responseBody != nil {
		length := int64(len(auditResponse.ResponseBodyBytes))
		responseLength = &length
	}

	// Get request body
	requestBody, err := NewRequestBody(auditRequest.Request)
	if err != nil {
		return nil, errors.Join(err, ErrInvalidRequestBody)
	}

	// Get audit attributes from request
	auditClaims, authenticationPrincipal, audiences, authenticationInfo, err :=
		AuditAttributesFromAuthorizationHeader(auditRequest.Request)
	if err != nil {
		return nil, err
	}

	// Get request scheme (http, https)
	scheme := auditRequest.Request.Scheme

	// Initialize authorization info if available
	var authorizationInfo []*auditV1.AuthorizationInfo
	if auditMetadata.AuditPermission != nil && auditMetadata.AuditPermissionGranted != nil {
		authorizationInfo = []*auditV1.AuthorizationInfo{
			NewAuthorizationInfo(
				auditMetadata.AuditResourceName,
				*auditMetadata.AuditPermission,
				*auditMetadata.AuditPermissionGranted)}
	}

	// Initialize labels if available
	var labels map[string]string
	if auditMetadata.AuditLabels != nil {
		labels = *auditMetadata.AuditLabels
	}

	// Initialize metadata/details
	var metadata *structpb.Struct
	if eventMetadata != nil {
		metadataStruct, err := structpb.NewStruct(eventMetadata)
		if err != nil {
			return nil, err
		}
		metadata = metadataStruct
	}

	// Get request and audit time
	var concreteRequestTime = time.Now().UTC()
	if auditRequest.RequestTime != nil {
		concreteRequestTime = *auditRequest.RequestTime
	}
	var concreteAuditTime = concreteRequestTime
	if auditMetadata.AuditTime != nil {
		concreteAuditTime = *auditMetadata.AuditTime
	}
	var concreteResponseTime = concreteRequestTime
	if auditResponse.ResponseTime != nil {
		concreteResponseTime = *auditResponse.ResponseTime
	}

	// Initialize the audit log entry
	event := auditV1.AuditLogEntry{
		LogName: auditMetadata.AuditLogName,
		ProtoPayload: &auditV1.AuditLog{
			ServiceName:        auditMetadata.AuditServiceName,
			OperationName:      auditMetadata.AuditOperationName,
			ResourceName:       auditMetadata.AuditResourceName,
			AuthenticationInfo: authenticationInfo,
			AuthorizationInfo:  authorizationInfo,
			RequestMetadata: NewRequestMetadata(
				auditRequest.Request,
				filteredRequestHeaders,
				auditRequest.RequestId,
				scheme,
				concreteRequestTime,
				auditRequest.RequestClientIP,
				authenticationPrincipal,
				audiences,
				auditClaims),
			Request: requestBody,
			ResponseMetadata: NewResponseMetadata(
				auditResponse.ResponseStatusCode,
				auditResponse.ResponseNumItems,
				responseLength,
				filteredResponseHeaders,
				concreteResponseTime),
			Response: responseBody,
			Metadata: metadata,
		},
		InsertId:      auditMetadata.AuditInsertId,
		Labels:        labels,
		CorrelationId: auditRequest.RequestCorrelationId,
		Timestamp:     timestamppb.New(concreteAuditTime),
		Severity:      auditMetadata.AuditLogSeverity,
	}
	return &event, nil
}

// NewPbInt64Value returns protobuf int64 wrapper if value is not nil.
func NewPbInt64Value(value *int64) *wrapperspb.Int64Value {
	if value != nil {
		return wrapperspb.Int64(*value)
	}
	return nil
}

// NewRequestMetadata returns initialized protobuf RequestMetadata object.
func NewRequestMetadata(
	request *pkgAuditCommon.ApiRequest,
	requestHeaders map[string]string,
	requestId *string,
	requestScheme string,
	requestTime time.Time,
	clientIp string,
	authenticationPrincipal string,
	audiences []string,
	auditClaims *structpb.Struct,
) *auditV1.RequestMetadata {
	agent := requestHeaders["User-Agent"]
	if agent == "" {
		agent = requestHeaders["user-agent"]
	}
	return &auditV1.RequestMetadata{
		CallerIp:                clientIp,
		CallerSuppliedUserAgent: agent,
		RequestAttributes: NewRequestAttributes(
			request,
			requestHeaders,
			requestId,
			requestScheme,
			requestTime,
			authenticationPrincipal,
			audiences,
			auditClaims,
		),
	}
}

// NewRequestAttributes returns initialized protobuf AttributeContext_Request object.
func NewRequestAttributes(
	request *pkgAuditCommon.ApiRequest,
	requestHeaders map[string]string,
	requestId *string,
	requestScheme string,
	requestTime time.Time,
	authenticationPrincipal string,
	audiences []string,
	auditClaims *structpb.Struct,
) *auditV1.AttributeContext_Request {

	rawQuery := request.URL.RawQuery
	var query *string
	if rawQuery != nil && *rawQuery != "" {
		escapedQuery := url.QueryEscape(*rawQuery)
		query = &escapedQuery
	}

	return &auditV1.AttributeContext_Request{
		Id:       requestId,
		Method:   pkgAuditCommon.StringToHttpMethod(request.Method),
		Headers:  requestHeaders,
		Path:     request.URL.Path,
		Host:     request.Host,
		Scheme:   requestScheme,
		Query:    query,
		Time:     timestamppb.New(requestTime),
		Protocol: request.Proto,
		Auth: &auditV1.AttributeContext_Auth{
			Principal: authenticationPrincipal,
			Audiences: audiences,
			Claims:    auditClaims,
		},
	}
}

// NewAuthorizationInfo returns protobuf AuthorizationInfo for the given parameters.
func NewAuthorizationInfo(resourceName, permission string, granted bool) *auditV1.AuthorizationInfo {
	return &auditV1.AuthorizationInfo{
		Resource:   resourceName,
		Permission: &permission,
		Granted:    &granted,
	}
}

// NewInsertId returns a correctly formatted insert id.
func NewInsertId(insertTime time.Time, location, workerId string, eventSequenceNumber uint64) string {
	return fmt.Sprintf("%d/%s/%s/%d", insertTime.UnixNano(), location, workerId, eventSequenceNumber)
}

// NewResponseMetadata returns protobuf response status with status code and short message.
func NewResponseMetadata(statusCode int, numResponseItems, responseSize *int64, headers map[string]string, responseTime time.Time) *auditV1.ResponseMetadata {

	var message *string
	if statusCode >= 400 && statusCode < 500 {
		text := "Client error"
		message = &text
	} else if statusCode >= 500 {
		text := "Server error"
		message = &text
	}

	var size *wrapperspb.Int64Value
	if responseSize != nil {
		size = wrapperspb.Int64(*responseSize)
	}
	return &auditV1.ResponseMetadata{
		StatusCode:   wrapperspb.Int32(int32(statusCode)),
		ErrorMessage: message,
		ErrorDetails: nil,
		ResponseAttributes: &auditV1.AttributeContext_Response{
			NumResponseItems: NewPbInt64Value(numResponseItems),
			Size:             size,
			Headers:          headers,
			Time:             timestamppb.New(responseTime),
		},
	}
}

// NewResponseBody converts the JSON byte response into a protobuf struct.
func NewResponseBody(response []byte) (*structpb.Struct, error) {

	// Return if nil
	if len(response) == 0 {
		return nil, nil
	}

	// Convert to protobuf struct
	return byteArrayToPbStruct(response)
}

// NewRequestBody converts the request body into a protobuf struct.
func NewRequestBody(request *pkgAuditCommon.ApiRequest) (*structpb.Struct, error) {

	if len(request.Body) == 0 {
		return nil, nil
	}

	// Convert to protobuf struct
	return byteArrayToPbStruct(request.Body)
}

// byteArrayToPbStruct converts a given json byte array into a protobuf struct.
func byteArrayToPbStruct(bytes []byte) (*structpb.Struct, error) {
	var bodyMap map[string]interface{}
	err := json.Unmarshal(bytes, &bodyMap)
	if err != nil {
		return nil, err
	}

	return structpb.NewStruct(bodyMap)
}

// FilterAndMergeHeaders filters ":authority", "Authorization", "B3" and "Host" headers as well as
// all headers starting with the prefixes "X-", "STACKIT-" and "grpcgateway-".
// Headers are merged if there is more than one value for a given name.
func FilterAndMergeHeaders(headers map[string][]string) map[string]string {
	var resultMap = make(map[string]string)
	skipHeaders := []string{":authority", "authorization", "b3", "host"}
	skipPrefixHeaders := []string{"x-", "stackit-", "grpcgateway-"}

	if len(headers) == 0 {
		return nil
	}

	for headerName, headerValues := range headers {
		headerLower := strings.ToLower(headerName)

		// Check if headers with a specific prefix is found
		skip := false
		for _, skipPrefix := range skipPrefixHeaders {
			if strings.HasPrefix(headerLower, skipPrefix) {
				skip = true
				break
			}
		}

		// Keep header if not on filter list or value is empty
		if !skip && !slices.Contains(skipHeaders, headerLower) && len(headerValues) > 0 {
			resultMap[headerName] = strings.Join(headerValues, ",")
		}
	}

	return resultMap
}

// NewAuditRoutingIdentifier instantiates a new auditApi.RoutableIdentifier for
// the given object ID and object type.
func NewAuditRoutingIdentifier(objectId string, objectType pkgAuditCommon.ObjectType) *pkgAuditCommon.RoutableIdentifier {
	return &pkgAuditCommon.RoutableIdentifier{
		Identifier: objectId,
		Type:       objectType,
	}
}

// AuditAttributesFromAuthorizationHeader extracts the following claims from given http.Request:
//   - auditClaims - filtered list of claims
//   - authenticationPrincipal - principal identifier
//   - audiences - list of audience claims
//   - authenticationInfo - information about the user or service-account authentication
func AuditAttributesFromAuthorizationHeader(request *pkgAuditCommon.ApiRequest) (
	*structpb.Struct,
	string,
	[]string,
	*auditV1.AuthenticationInfo,
	error,
) {

	var authenticationPrincipal = "none/none"
	var principalId = "none"
	var principalEmail *string
	emptyClaims, err := structpb.NewStruct(make(map[string]interface{}))
	if err != nil {
		return nil, authenticationPrincipal, nil, nil, err
	}
	var auditClaims = emptyClaims
	var serviceAccountName *string
	audiences := make([]string, 0)
	var delegationInfo []*auditV1.ServiceAccountDelegationInfo

	authorizationHeaders := request.Header["Authorization"]
	if len(authorizationHeaders) == 0 {
		// fallback for grpc where headers/metadata keys are lowercase
		authorizationHeaders = request.Header["authorization"]
	}
	authorizationHeader := strings.Join(authorizationHeaders, ",")
	trimmedAuthorizationHeader := strings.TrimSpace(authorizationHeader)
	if trimmedAuthorizationHeader != "" {

		// Parse claims
		token, err := parseToken(trimmedAuthorizationHeader)
		if err != nil {
			return nil, authenticationPrincipal, nil, nil, err
		}

		filteredClaims, err := parseClaimsFromAuthorizationHeader(token)
		if err != nil {
			return nil, authenticationPrincipal, nil, nil, err
		}

		// Convert filtered claims to protobuf struct
		auditClaimsStruct, err := structpb.NewStruct(filteredClaims)
		if err != nil {
			return nil, authenticationPrincipal, nil, nil, err
		}
		auditClaims = auditClaimsStruct

		// Extract principal data
		authenticationPrincipal = extractAuthenticationPrincipal(token)
		principalId, principalEmail = extractSubjectAndEmail(token)

		// Extract service account delegation info data
		actClaim, hasActClaim := token.Get("act")
		if hasActClaim {
			actMap := map[string]interface{}{"act": actClaim}
			delegationInfo = extractServiceAccountDelegationInfo(actMap)
		}

		// Extract audiences data
		audiences = token.Audience()

		// Extract project id and service account id
		projectId := getTokenClaim(token, TokenClaimStackitProjectId)
		serviceAccountId := getTokenClaim(token, TokenClaimStackitServiceAccountId)

		// Calculate service account name if project and service account ids are available
		if projectId != nil && serviceAccountId != nil {
			accountName := fmt.Sprintf("projects/%s/service-accounts/%s", *projectId, *serviceAccountId)
			serviceAccountName = &accountName
		}
	}

	authenticationInfo := auditV1.AuthenticationInfo{
		PrincipalId:                  principalId,
		PrincipalEmail:               principalEmail,
		ServiceAccountName:           serviceAccountName,
		ServiceAccountDelegationInfo: delegationInfo,
	}

	return auditClaims, authenticationPrincipal, audiences, &authenticationInfo, nil
}

func getTokenClaim(token jwt.Token, claimName string) *string {
	claim, claimExists := token.Get(claimName)
	if claimExists {
		claimString := fmt.Sprintf("%s", claim)
		return &claimString
	}
	return nil
}

func extractAuthenticationPrincipal(token jwt.Token) string {
	subject := token.Subject()
	issuer := token.Issuer()

	var principal = "none/none"
	if subject != "" && issuer != "" {
		principal = fmt.Sprintf("%s/%s", url.QueryEscape(subject), url.QueryEscape(issuer))
	}
	return principal
}

func parseToken(authorizationHeader string) (jwt.Token, error) {
	parts := strings.Split(authorizationHeader, " ")
	if len(parts) != 2 {
		return nil, ErrInvalidAuthorizationHeaderValue
	}
	if !strings.EqualFold(parts[0], "Bearer") {
		return nil, ErrTokenIsNotBearerToken
	}
	jwtString := parts[1]
	authorizationHeaderParts := strings.Split(jwtString, ".")

	if len(authorizationHeaderParts) == 3 {
		token, err := jwt.ParseInsecure([]byte(parts[1]))
		if err != nil {
			return nil, ErrInvalidBearerToken
		}
		return token, nil
	}
	return nil, ErrInvalidBearerToken
}

func parseClaimsFromAuthorizationHeader(token jwt.Token) (map[string]interface{}, error) {
	claimsMap, err := token.AsMap(context.Background())
	if err != nil {
		return nil, err
	}

	claims := map[string]interface{}{}
	if len(token.Audience()) > 0 {
		var audiences []any
		for _, audience := range token.Audience() {
			audiences = append(audiences, audience)
		}
		claims["aud"] = audiences
	}

	for key := range claimsMap {
		if key == "aud" {
			continue
		}
		value := claimsMap[key]
		t, isTime := value.(time.Time)
		if isTime {
			claims[key] = t.String()
		} else if value != nil && value != "" {
			claims[key] = value
		}
	}
	return claims, nil
}

func extractServiceAccountDelegationInfoDetails(actClaims map[string]interface{}) []*auditV1.ServiceAccountDelegationInfo {
	principalId, principalEmail := extractSubjectAndEmailFromActClaims(actClaims)

	delegation := auditV1.ServiceAccountDelegationInfo{Authority: &auditV1.ServiceAccountDelegationInfo_IdpPrincipal_{IdpPrincipal: &auditV1.ServiceAccountDelegationInfo_IdpPrincipal{
		PrincipalId:     principalId,
		PrincipalEmail:  principalEmail,
		ServiceMetadata: nil,
	}}}

	delegations := []*auditV1.ServiceAccountDelegationInfo{&delegation}
	nestedDelegations := extractServiceAccountDelegationInfo(actClaims)
	if len(nestedDelegations) > 0 {
		return append(delegations, nestedDelegations...)
	}
	return delegations
}

func extractServiceAccountDelegationInfo(claims map[string]interface{}) []*auditV1.ServiceAccountDelegationInfo {
	actor, hasActor := claims["act"]
	if hasActor {
		actorMap, hasActorClaim := actor.(map[string]interface{})
		if hasActorClaim {
			return extractServiceAccountDelegationInfoDetails(actorMap)
		}
	}
	return nil
}

func extractSubjectAndEmailFromActClaims(actClaim map[string]interface{}) (string, string) {
	var principalEmail string
	principalId := fmt.Sprintf("%s", actClaim["sub"])
	principalEmailRaw := actClaim["email"]
	if principalEmailRaw == nil {
		principalEmail = EmailAddressDoNotReplyAtStackItDotCloud
	} else {
		principalEmail = fmt.Sprintf("%s", principalEmailRaw)
	}
	return principalId, principalEmail
}

func extractSubjectAndEmail(token jwt.Token) (string, *string) {
	var principalEmail *string
	principalId := token.Subject()
	emailClaim, hasEmail := token.Get("email")
	if hasEmail {
		trimmedEmail := strings.TrimSpace(fmt.Sprintf("%s", emailClaim))
		if trimmedEmail != "" {
			principalEmail = &trimmedEmail
		}
	}
	return principalId, principalEmail
}