From d5ab35a444523baa0586fcb9513d6ae4f2518946 Mon Sep 17 00:00:00 2001
From: aliernfrog <45766489+aliernfrog@users.noreply.github.com>
Date: Thu, 28 Sep 2023 04:03:07 +0300
Subject: [PATCH] feat(YouTube - Premium heading): Allow using default heading
 (#3029)

Co-authored-by: oSumAtrIX <johan.melkonyan1@web.de>
---
 .../branding/header/PremiumHeadingPatch.kt    | 56 ++++++++++++-------
 1 file changed, 37 insertions(+), 19 deletions(-)

diff --git a/src/main/kotlin/app/revanced/patches/youtube/layout/branding/header/PremiumHeadingPatch.kt b/src/main/kotlin/app/revanced/patches/youtube/layout/branding/header/PremiumHeadingPatch.kt
index 7f9db9c2a..48fc6d9b5 100644
--- a/src/main/kotlin/app/revanced/patches/youtube/layout/branding/header/PremiumHeadingPatch.kt
+++ b/src/main/kotlin/app/revanced/patches/youtube/layout/branding/header/PremiumHeadingPatch.kt
@@ -5,39 +5,57 @@ import app.revanced.patcher.patch.PatchException
 import app.revanced.patcher.patch.ResourcePatch
 import app.revanced.patcher.patch.annotation.CompatiblePackage
 import app.revanced.patcher.patch.annotation.Patch
-import java.nio.file.Files
-import java.nio.file.StandardCopyOption
-import kotlin.io.path.exists
+import app.revanced.patcher.patch.options.types.BooleanPatchOption.Companion.booleanPatchOption
+import kotlin.io.path.copyTo
 
 @Patch(
     name = "Premium heading",
-    description = "Shows premium branding on the home screen.",
+    description = "Show or hide the premium heading.",
     compatiblePackages = [
         CompatiblePackage("com.google.android.youtube")
     ]
 )
 @Suppress("unused")
 object PremiumHeadingPatch : ResourcePatch() {
+    private const val DEFAULT_HEADING_RES = "yt_wordmark_header"
+    private const val PREMIUM_HEADING_RES = "yt_premium_wordmark_header"
+
+    private val usePremiumHeading by booleanPatchOption(
+        key = "usePremiumHeading",
+        default = true,
+        title = "Use premium heading",
+        description = "Whether to use the premium heading.",
+        required = true,
+    )
+
     override fun execute(context: ResourceContext) {
         val resDirectory = context["res"]
-        if (!resDirectory.isDirectory) throw PatchException("The res folder can not be found.")
 
-        val (original, replacement) = "yt_premium_wordmark_header" to "yt_wordmark_header"
-        val modes = arrayOf("light", "dark")
+        val (original, replacement) = if (usePremiumHeading!!)
+            DEFAULT_HEADING_RES to PREMIUM_HEADING_RES
+        else
+            PREMIUM_HEADING_RES to DEFAULT_HEADING_RES
 
-        arrayOf("xxxhdpi", "xxhdpi", "xhdpi", "hdpi", "mdpi").forEach { size ->
-            val headingDirectory = resDirectory.resolve("drawable-$size")
-            modes.forEach { mode ->
-                val fromPath = headingDirectory.resolve("${original}_$mode.png").toPath()
-                val toPath = headingDirectory.resolve("${replacement}_$mode.png").toPath()
+        val variants = arrayOf("light", "dark")
 
-                if (!fromPath.exists())
-                    throw PatchException("The file $fromPath does not exist in the resources. Therefore, this patch can not succeed.")
-                Files.copy(
-                    fromPath,
-                    toPath,
-                    StandardCopyOption.REPLACE_EXISTING
-                )
+        arrayOf(
+            "xxxhdpi",
+            "xxhdpi",
+            "xhdpi",
+            "hdpi",
+            "mdpi"
+        ).mapNotNull { dpi ->
+            resDirectory.resolve("drawable-$dpi").takeIf { it.exists() }?.toPath()
+        }.also {
+            if (it.isEmpty())
+                throw PatchException("The drawable folder can not be found. Therefore, the patch can not be applied.")
+        }.forEach { path ->
+
+            variants.forEach { mode ->
+                val fromPath = path.resolve("${original}_$mode.png")
+                val toPath = path.resolve("${replacement}_$mode.png")
+
+                fromPath.copyTo(toPath, true)
             }
         }
     }