class StytchB2B::Sessions

def authenticate(
  session_token: nil,
  session_duration_minutes: nil,
  session_jwt: nil,
  session_custom_claims: nil,
  authorization_check: nil
)
  headers = {}
  request = {}
  request[:session_token] = session_token unless session_token.nil?
  request[:session_duration_minutes] = session_duration_minutes unless session_duration_minutes.nil?
  request[:session_jwt] = session_jwt unless session_jwt.nil?
  request[:session_custom_claims] = session_custom_claims unless session_custom_claims.nil?
  request[:authorization_check] = authorization_check unless authorization_check.nil?
  post_request('/v1/b2b/sessions/authenticate', request, headers)
end

def authenticate_jwt(
  session_jwt,
  max_token_age_seconds: nil,
  session_duration_minutes: nil,
  session_custom_claims: nil,
  authorization_check: nil,
  clock_tolerance_seconds: nil
)
  max_token_age_seconds = 300 if max_token_age_seconds.nil?
  clock_tolerance_seconds = 0 if clock_tolerance_seconds.nil?
  if max_token_age_seconds == 0
    return authenticate(
      session_jwt: session_jwt,
      session_duration_minutes: session_duration_minutes,
      session_custom_claims: session_custom_claims,
      authorization_check: authorization_check
    )
  end
  decoded_jwt = authenticate_jwt_local(session_jwt, max_token_age_seconds: max_token_age_seconds, authorization_check: authorization_check, clock_tolerance_seconds: clock_tolerance_seconds)
  return decoded_jwt unless decoded_jwt.nil?
  authenticate(
    session_jwt: session_jwt,
    session_duration_minutes: session_duration_minutes,
    session_custom_claims: session_custom_claims,
    authorization_check: authorization_check
  )
rescue StandardError
  # JWT could not be verified locally. Check with the Stytch API.
  authenticate(
    session_jwt: session_jwt,
    session_duration_minutes: session_duration_minutes,
    session_custom_claims: session_custom_claims,
    authorization_check: authorization_check
  )
end

def authenticate_jwt_local(session_jwt, max_token_age_seconds: nil, authorization_check: nil, clock_tolerance_seconds: nil)
  max_token_age_seconds = 300 if max_token_age_seconds.nil?
  clock_tolerance_seconds = 0 if clock_tolerance_seconds.nil?
  issuer = 'stytch.com/' + @project_id
  begin
    decoded_token = JWT.decode session_jwt, nil, true,
                               { jwks: @jwks_loader, iss: issuer, verify_iss: true, aud: @project_id, verify_aud: true, algorithms: ['RS256'], nbf_leeway: clock_tolerance_seconds }
    session = decoded_token[0]
    iat_time = Time.at(session['iat']).to_datetime
    return nil unless iat_time + max_token_age_seconds >= Time.now
    session = marshal_jwt_into_session(session)
  rescue JWT::InvalidIssuerError
    raise Stytch::JWTInvalidIssuerError
  rescue JWT::InvalidAudError
    raise Stytch::JWTInvalidAudienceError
  rescue JWT::ExpiredSignature
    raise Stytch::JWTExpiredSignatureError
  rescue JWT::IncorrectAlgorithm
    raise Stytch::JWTIncorrectAlgorithmError
  end
  # Do the auth check - intentionally don't rescue errors from here
  if authorization_check && session['roles']
    @policy_cache.perform_authorization_check(
      subject_roles: session['roles'],
      subject_org_id: session['member_session']['organization_id'],
      authorization_check: authorization_check
    )
  end
  session
end

def exchange(
  organization_id:,
  session_token: nil,
  session_jwt: nil,
  session_duration_minutes: nil,
  session_custom_claims: nil,
  locale: nil
)
  headers = {}
  request = {
    organization_id: organization_id
  }
  request[:session_token] = session_token unless session_token.nil?
  request[:session_jwt] = session_jwt unless session_jwt.nil?
  request[:session_duration_minutes] = session_duration_minutes unless session_duration_minutes.nil?
  request[:session_custom_claims] = session_custom_claims unless session_custom_claims.nil?
  request[:locale] = locale unless locale.nil?
  post_request('/v1/b2b/sessions/exchange', request, headers)
end

def exchange_access_token(
  access_token:,
  session_duration_minutes: nil,
  session_custom_claims: nil
)
  headers = {}
  request = {
    access_token: access_token
  }
  request[:session_duration_minutes] = session_duration_minutes unless session_duration_minutes.nil?
  request[:session_custom_claims] = session_custom_claims unless session_custom_claims.nil?
  post_request('/v1/b2b/sessions/exchange_access_token', request, headers)
end

def get(
  organization_id:,
  member_id:
)
  headers = {}
  query_params = {
    organization_id: organization_id,
    member_id: member_id
  }
  request = request_with_query_params('/v1/b2b/sessions', query_params)
  get_request(request, headers)
end

def get_jwks(
  project_id:
)
  headers = {}
  query_params = {}
  request = request_with_query_params("/v1/b2b/sessions/jwks/#{project_id}", query_params)
  get_request(request, headers)
end

def initialize(connection, project_id, policy_cache)
  @connection = connection
  @policy_cache = policy_cache
  @project_id = project_id
  @cache_last_update = 0
  @jwks_loader = lambda do |options|
    @cached_keys = nil if options[:invalidate] && @cache_last_update < Time.now.to_i - 300
    @cached_keys ||= begin
      @cache_last_update = Time.now.to_i
      keys = []
      get_jwks(project_id: @project_id)['keys'].each do |r|
        keys << r
      end
      { keys: keys }
    end
  end
end

def marshal_jwt_into_session(jwt)
  stytch_claim = 'https://stytch.com/session'
  organization_claim = 'https://stytch.com/organization'
  expires_at = jwt[stytch_claim]['expires_at'] || Time.at(jwt['exp']).to_datetime.utc.strftime('%Y-%m-%dT%H:%M:%SZ')
  # The custom claim set is all the claims in the payload except for the standard claims and
  # the Stytch session claim. The cleanest way to collect those seems to be naming what we want
  # to omit and filtering the rest to collect the custom claims.
  reserved_claims = ['aud', 'exp', 'iat', 'iss', 'jti', 'nbf', 'sub', stytch_claim, organization_claim]
  custom_claims = jwt.reject { |key, _| reserved_claims.include?(key) }
  {
    'member_session' => {
      'session_id' => jwt[stytch_claim]['id'],
      'organization_id' => jwt[organization_claim]['organization_id'],
      'member_id' => jwt['sub'],
      # DEPRECATED: Use member_id instead
      'user_id' => jwt['sub'],
      'started_at' => jwt[stytch_claim]['started_at'],
      'last_accessed_at' => jwt[stytch_claim]['last_accessed_at'],
      # For JWTs that include it, prefer the inner expires_at claim.
      'expires_at' => expires_at,
      'attributes' => jwt[stytch_claim]['attributes'],
      'authentication_factors' => jwt[stytch_claim]['authentication_factors'],
      'custom_claims' => custom_claims
    },
    'roles' => jwt[stytch_claim]['roles']
  }
end

def migrate(
  session_token:,
  organization_id:,
  session_duration_minutes: nil,
  session_custom_claims: nil
)
  headers = {}
  request = {
    session_token: session_token,
    organization_id: organization_id
  }
  request[:session_duration_minutes] = session_duration_minutes unless session_duration_minutes.nil?
  request[:session_custom_claims] = session_custom_claims unless session_custom_claims.nil?
  post_request('/v1/b2b/sessions/migrate', request, headers)
end

def revoke(
  member_session_id: nil,
  session_token: nil,
  session_jwt: nil,
  member_id: nil,
  method_options: nil
)
  headers = {}
  headers = headers.merge(method_options.to_headers) unless method_options.nil?
  request = {}
  request[:member_session_id] = member_session_id unless member_session_id.nil?
  request[:session_token] = session_token unless session_token.nil?
  request[:session_jwt] = session_jwt unless session_jwt.nil?
  request[:member_id] = member_id unless member_id.nil?
  post_request('/v1/b2b/sessions/revoke', request, headers)
end