Keep the screen on while playing in the offline player

app/src/main/java/com/github/libretube/ui/activities/MainActivity.kt

@@ -167,7 +167,7 @@ class MainActivity : BaseActivity() {
                 }
                 R.id.searchResultFragment -> {
                     navController.popBackStack(R.id.searchFragment, true) ||
-                            navController.popBackStack()
+                        navController.popBackStack()
                 }
                 else -> {
                     navController.popBackStack()

app/src/main/java/com/github/libretube/ui/fragments/PlayerFragment.kt

@@ -919,11 +919,6 @@ class PlayerFragment : BaseFragment(), OnlinePlayerOptions {
             }
 
             override fun onPlaybackStateChanged(playbackState: Int) {
-                exoPlayerView.keepScreenOn = !(
-                    playbackState == Player.STATE_IDLE ||
-                        playbackState == Player.STATE_ENDED
-                    )
-
                 // save the watch position to the database
                 // only called when the position is unequal to 0, otherwise it would become reset
                 // before the player can seek to the saved position from videos of the queue

app/src/main/java/com/github/libretube/ui/views/CustomExoPlayerView.kt

@@ -177,6 +177,15 @@ internal class CustomExoPlayerView(
                     updatePlayPauseButton()
                 }
             }
+
+            override fun onPlaybackStateChanged(playbackState: Int) {
+                // keep the screen on if the video is not ended or paused
+                keepScreenOn = !(
+                    listOf(Player.STATE_IDLE, Player.STATE_ENDED)
+                        .contains(playbackState)
+                    )
+                super.onPlaybackStateChanged(playbackState)
+            }
         })
 
         playerViewModel?.isFullscreen?.observe(viewLifecycleOwner!!) { isFullscreen ->