diff --git a/main.cpp b/main.cpp
index c3a72c6..f621e2c 100644
--- a/main.cpp
+++ b/main.cpp
@@ -1,25 +1,20 @@
 #include "config.hpp"
-#include "reader.hpp"
 #include "record.hpp"
+#include "writer.hpp"
 #include <fstream>
 #include <iostream>
 #include <sys/stat.h>
 #include <sys/types.h>
 
-int generate_file(int N, const std::string& filename)
+void generate_file(int N, const std::string& filename)
 {
-    std::ofstream out(filename);
-    if (!out) {
-        std::cerr << "Failed to open output file: " << filename << '\n';
-        return 1;
-    }
+    std::ofstream out_stream(filename);
+    Writer output(out_stream);
 
     for (int i = 0; i < N; i++) {
         Record rec = Record::random();
-        out << rec << '\n';
+        output.write(rec);
     }
-
-    return 0;
 }
 
 int sort_file(const Configuration& opts)
@@ -35,10 +30,11 @@ int main(int argc, char** argv)
         Configuration opts = Configuration::parse_args(argc, argv);
 
         if (opts.generate_data) {
-            return generate_file(opts.N, opts.output_file);
+            generate_file(opts.N, opts.output_file);
+            return 0;
         }
 
-        return sort_file(opts);
+        sort_file(opts);
 
     } catch (const std::exception& e) {
         std::cerr << "Error: " << e.what() << '\n';
diff --git a/writer.hpp b/writer.hpp
index 99dbb84..9fe4e80 100644
--- a/writer.hpp
+++ b/writer.hpp
@@ -9,19 +9,21 @@ private:
     std::ostream& output_stream;
     size_t buffer_size;
     size_t total_writes_count;
-    size_t current_idx; // Current position in buffer
+    size_t current_index;
 
     // returns true on success
     // doesn't really flush for performace reasons
-    // could be changed to really flush to monitor the impac of amount of writes
+    // could be changed to really flush to monitor the impact of amount of writes
     bool flush_buffer()
     {
-        for (size_t i = 0; i < current_idx; ++i) {
+        for (size_t i = 0; i < current_index; ++i) {
             if (!(output_stream << write_buffer[i])) {
                 return false;
             }
         }
 
+        current_index = 0;
+
         return true;
     }
 
@@ -31,17 +33,22 @@ public:
         , output_stream(output_stream)
         , buffer_size(buffer_size)
         , total_writes_count(0)
-        , current_idx(0)
+        , current_index(0)
     {
     }
 
+    ~Writer()
+    {
+        flush_buffer();
+    }
+
     bool write(const Record& record)
     {
-        write_buffer[current_idx] = record;
-        current_idx++;
-        total_writes_count++;
+        write_buffer[current_index] = record;
+        ++current_index;
+        ++total_writes_count;
 
-        if (current_idx >= buffer_size) {
+        if (current_index >= buffer_size) {
             return flush_buffer();
         }