diff --git a/build.gradle b/build.gradle
index 91d1ea5..8b59a9b 100644
--- a/build.gradle
+++ b/build.gradle
@@ -17,6 +17,7 @@ plugins {
     id 'com.google.protobuf' version '0.8.19'
 
     id 'eclipse'
+    id 'idea'
 }
 
 compileJava.options.encoding = "UTF-8"
@@ -25,38 +26,42 @@ compileTestJava.options.encoding = "UTF-8"
 sourceCompatibility = JavaVersion.VERSION_17
 targetCompatibility = JavaVersion.VERSION_17
 
+var shouldGenerateProto = System.getenv("GENERATE_PROTO") == "true"
+System.out.println(shouldGenerateProto ?
+        "Generating proto files" :
+        "Skipping proto generation")
+
 repositories {
     mavenCentral()
     jcenter()
 }
 
 protobuf {
-	protoc {
+    protoc {
         artifact = 'com.google.protobuf:protoc:3.24.3'
     }
-	plugins {
-		quickbuf {
-			artifact = 'us.hebi.quickbuf:protoc-gen-quickbuf:1.3.1'
-		}
-	}
 
-    if (System.getenv("GENERATE_PROTO") == "true") {
-        generateProtoTasks {
-            all().each { task ->
-                task.builtins {
-                    remove java
-                }
-                task.plugins {
-                    quickbuf {
-                        //option 'store_unknown_fields=true'
-                        outputSubDir = ''
-                    }
+    plugins {
+        quickbuf {
+            artifact = 'us.hebi.quickbuf:protoc-gen-quickbuf:1.3.1'
+        }
+    }
+
+    generateProtoTasks {
+        all().each { task ->
+            task.builtins {
+                remove java
+            }
+            task.plugins {
+                quickbuf {
+                    //option 'store_unknown_fields=true'
+                    outputSubDir = ''
                 }
             }
         }
     }
 
-	generatedFilesBaseDir = "$projectDir/src/generated/"
+    generatedFilesBaseDir = "$projectDir/src/generated/"
 }
 
 dependencies {
@@ -143,6 +148,15 @@ sourceSets {
     }
 }
 
+idea {
+    module {
+        // proto files and generated Java files are automatically added as
+        // source dirs.
+        // If you have additional sources, add them here:
+        sourceDirs += file('/proto/')
+    }
+}
+
 eclipse {
     classpath {
         file.whenMerged { cp ->
@@ -172,7 +186,11 @@ public final class BuildConfig {
 task updateProto {
     group 'Build'
     description 'Cleans generated proto folder and regenerates protos'
-    dependsOn 'clean'
+
+    if (shouldGenerateProto) {
+        dependsOn 'clean'
+    }
+
     dependsOn 'generateProto'
 }