diff --git a/src/main/kotlin/Diff.kt b/src/main/kotlin/Diff.kt
index 9b3250f..f4293ca 100644
--- a/src/main/kotlin/Diff.kt
+++ b/src/main/kotlin/Diff.kt
@@ -14,13 +14,12 @@ import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.text.font.FontFamily
 import androidx.compose.ui.unit.dp
-import org.eclipse.jgit.diff.DiffEntry
 import theme.primaryTextColor
 
 @Composable
-fun Diff(gitManager: GitManager, diffEntry: DiffEntry, onCloseDiffView: () -> Unit) {
-    val text = remember(diffEntry) {
-        gitManager.diffFormat(diffEntry)
+fun Diff(gitManager: GitManager, diffEntryType: DiffEntryType, onCloseDiffView: () -> Unit) {
+    val text = remember(diffEntryType.diffEntry) {
+        gitManager.diffFormat(diffEntryType)
     }
 
     Card(
diff --git a/src/main/kotlin/DiffEntryType.kt b/src/main/kotlin/DiffEntryType.kt
new file mode 100644
index 0000000..5b67333
--- /dev/null
+++ b/src/main/kotlin/DiffEntryType.kt
@@ -0,0 +1,7 @@
+import org.eclipse.jgit.diff.DiffEntry
+
+sealed class DiffEntryType(val diffEntry: DiffEntry) {
+    class CommitDiff(diffEntry: DiffEntry): DiffEntryType(diffEntry)
+    class UnstagedDiff(diffEntry: DiffEntry): DiffEntryType(diffEntry)
+    class StagedDiff(diffEntry: DiffEntry): DiffEntryType(diffEntry)
+}
diff --git a/src/main/kotlin/GitManager.kt b/src/main/kotlin/GitManager.kt
index 5187535..017059f 100644
--- a/src/main/kotlin/GitManager.kt
+++ b/src/main/kotlin/GitManager.kt
@@ -126,7 +126,8 @@ class GitManager {
     val hasUncommitedChanges: StateFlow<Boolean>
         get() = statusManager.hasUncommitedChanges
 
-    fun diffFormat(diffEntry: DiffEntry): String {
+    fun diffFormat(diffEntryType: DiffEntryType): String {
+        val diffEntry = diffEntryType.diffEntry
         val byteArrayOutputStream = ByteArrayOutputStream()
 
         DiffFormatter(byteArrayOutputStream).use { formatter ->
@@ -136,9 +137,9 @@ class GitManager {
             val oldTree = DirCacheIterator(repo.readDirCache())
             val newTree = FileTreeIterator(repo)
 
-            println(diffEntry)
-            formatter.scan(oldTree, newTree) //TODO Should only be set when using diff for unstaged changes
-//            formatter.format(oldTree, newTree)
+            if(diffEntryType is DiffEntryType.UnstagedDiff)
+                formatter.scan(oldTree, newTree)
+
             formatter.format(diffEntry)
             formatter.flush()
         }
diff --git a/src/main/kotlin/RepositorySelected.kt b/src/main/kotlin/RepositorySelected.kt
index 0b52520..b38fe21 100644
--- a/src/main/kotlin/RepositorySelected.kt
+++ b/src/main/kotlin/RepositorySelected.kt
@@ -20,7 +20,7 @@ fun RepositorySelected(gitManager: GitManager, repository: Repository) {
     }
 
     var diffSelected by remember {
-        mutableStateOf<DiffEntry?>(null)
+        mutableStateOf<DiffEntryType?>(null)
     }
     var uncommitedChangesSelected by remember {
         mutableStateOf<Boolean>(false)
@@ -79,7 +79,7 @@ fun RepositorySelected(gitManager: GitManager, repository: Repository) {
                     else -> {
                         Diff(
                             gitManager = gitManager,
-                            diffEntry = diffEntry,
+                            diffEntryType = diffEntry,
                             onCloseDiffView = { diffSelected = null })
                     }
                 }
@@ -94,9 +94,11 @@ fun RepositorySelected(gitManager: GitManager, repository: Repository) {
             if (uncommitedChangesSelected) {
                 UncommitedChanges(
                     gitManager = gitManager,
-                    onDiffEntrySelected = { diffEntry ->
-                        println(diffEntry.filePath)
-                        diffSelected = diffEntry
+                    onStagedDiffEntrySelected = { diffEntry ->
+                        diffSelected = DiffEntryType.StagedDiff(diffEntry)
+                    },
+                    onUnstagedDiffEntrySelected = { diffEntry ->
+                        diffSelected = DiffEntryType.UnstagedDiff(diffEntry)
                     }
                 )
             } else {
@@ -104,7 +106,7 @@ fun RepositorySelected(gitManager: GitManager, repository: Repository) {
                     CommitChanges(
                         commitDiff = it,
                         onDiffSelected = { diffEntry ->
-                            diffSelected = diffEntry
+                            diffSelected = DiffEntryType.CommitDiff(diffEntry)
                         }
                     )
                 }
diff --git a/src/main/kotlin/UncommitedChanges.kt b/src/main/kotlin/UncommitedChanges.kt
index 52c2cca..d9d1072 100644
--- a/src/main/kotlin/UncommitedChanges.kt
+++ b/src/main/kotlin/UncommitedChanges.kt
@@ -33,7 +33,8 @@ import theme.headerBackground
 @Composable
 fun UncommitedChanges(
     gitManager: GitManager,
-    onDiffEntrySelected: (DiffEntry) -> Unit,
+    onStagedDiffEntrySelected: (DiffEntry) -> Unit,
+    onUnstagedDiffEntrySelected: (DiffEntry) -> Unit,
 ) {
     val stageStatusState = gitManager.stageStatus.collectAsState()
     val stageStatus = stageStatusState.value
@@ -64,7 +65,7 @@ fun UncommitedChanges(
             title = "Staged",
             optionIcon = Icons.Default.Close,
             diffEntries = staged,
-            onDiffEntrySelected = onDiffEntrySelected,
+            onDiffEntrySelected = onStagedDiffEntrySelected,
             onDiffEntryOptionSelected = {
                 gitManager.unstage(it)
             }
@@ -78,7 +79,7 @@ fun UncommitedChanges(
             title = "Unstaged",
             optionIcon = Icons.Default.Add,
             diffEntries = unstaged,
-            onDiffEntrySelected = onDiffEntrySelected,
+            onDiffEntrySelected = onUnstagedDiffEntrySelected,
             onDiffEntryOptionSelected = {
                 gitManager.stage(it)
             }