diff --git a/app/src/main/java/com/clemhaowen/dm_td2/task/TaskActivity.kt b/app/src/main/java/com/clemhaowen/dm_td2/task/TaskActivity.kt
index 6f98cce0a76444f3b0b53451f498fe03831fc386..cb945c986a04160d72bb204f9e78cb1c79b93f45 100644
--- a/app/src/main/java/com/clemhaowen/dm_td2/task/TaskActivity.kt
+++ b/app/src/main/java/com/clemhaowen/dm_td2/task/TaskActivity.kt
@@ -21,23 +21,12 @@ class TaskActivity : AppCompatActivity() {
         var resultCode = 666
         val titleEditText = findViewById<EditText>(R.id.titleEditText)
         val descriptionEditText = findViewById<EditText>(R.id.descriptionEditText)
-        val serializableExtra = intent?.getSerializableExtra("editTask")
-        var task : Task
-        var id : String = ""
-        if (serializableExtra != null){
-            task = serializableExtra as Task
-            titleEditText.setText(task.title)
-            descriptionEditText.setText(task.description)
-            resultCode = 667
-            id = task.id
-        }
+        //val serializableExtra = intent?.getSerializableExtra("editTask")
+        val task = intent?.getSerializableExtra("editTask") as? Task
+        titleEditText.setText(task?.title)
+        descriptionEditText.setText(task?.description)
         createButton.setOnClickListener{
-            var newTask : Task
-
-            if (resultCode == 666){
-                id = java.util.UUID.randomUUID().toString()
-            }
-            newTask = Task(id = id, title = titleEditText.text.toString(), description = descriptionEditText.text.toString())
+            var newTask = Task(id = task?.id ?: java.util.UUID.randomUUID().toString(), title = titleEditText.text.toString(), description = descriptionEditText.text.toString())
 
             intent.putExtra("newTask", newTask)
             setResult(resultCode, intent)