diff --git a/main.py b/main.py
index c01e4d1..b565d07 100644
--- a/main.py
+++ b/main.py
@@ -1,4 +1,5 @@
 import src.task_splitter.code_processor as splitter
+import src.task_splitter.code_generator as generator
 
 input_file = "source-code/tasks.c"  # Replace this with your actual file name
 output_dir = "build/"
@@ -6,6 +7,9 @@ output_dir = "build/"
 functions = splitter.extract_functions_from_file(input_file)
 print(f"Found {len(functions)} functions in the file.")
 
+generator.build_code_from_functions(functions)
+print("Code has been generated for each function.")
+
 splitter.create_independent_files(functions, output_dir)
 print("Functions have been split into separate files with their own main functions.")
 
diff --git a/src/task_splitter/code_generator.py b/src/task_splitter/code_generator.py
new file mode 100644
index 0000000..8feabdb
--- /dev/null
+++ b/src/task_splitter/code_generator.py
@@ -0,0 +1,47 @@
+
+import ast
+
+def __get_include_from_annotation(libdeps):
+    
+    result = ""
+    
+    for lib in libdeps:
+        result += f"#include <{lib}>\n"
+        
+    return result
+
+def __get_main_function_from_function(fun):
+    
+    main_function = f"""
+int main() {{
+    
+    {fun['name']}({', '.join(['0' for _ in fun['args'].split(',')])});
+    return 0;
+}}\n
+    """
+    return main_function
+
+def get_code_from_function(fun):
+    
+    # Get the function content
+    function_content = f"{fun['return_type']} {fun['name']}({fun['args']}) {fun['content']}"
+    
+    # Get the main function
+    main_function = __get_main_function_from_function(fun)
+    
+    # Get the dependencies from the annotations
+    libdeps = ""
+    for ann in fun['annotations']:
+        if ann['tag'] == "@LibDeps":
+            libdeps= ast.literal_eval(ann['args'])
+            break
+    
+    # Get the include statements
+    include_statements = __get_include_from_annotation(libdeps=libdeps)
+    
+    return include_statements + '\n' + function_content + '\n' + main_function
+
+def build_code_from_functions(functions):
+    
+    for fun in functions:
+        fun["code"] = get_code_from_function(fun)
diff --git a/src/task_splitter/code_processor.py b/src/task_splitter/code_processor.py
index 6975b9f..ff20bf7 100644
--- a/src/task_splitter/code_processor.py
+++ b/src/task_splitter/code_processor.py
@@ -69,24 +69,11 @@ def create_independent_files(functions, output_dir):
     
     for fun in functions:
         
-        function_content = f"{fun['return_type']} {fun['name']}({fun['args']}) {fun['content']}"
-        
         filename = f"{fun['name']}.c"
         with open(output_dir + filename, "w") as f:
             
-            # Write the original function
-            f.write(function_content)
-            f.write("\n")
-            
-            # Add a main function to call the function
-            main_function = f"""
-int main() {{
-    // Assuming that the function has no return value
-    {fun['name']}({', '.join(['0' for _ in fun['args'].split(',')])});
-    return 0;
-}}
-            """
-            f.write(main_function)
+            # Write the code
+            f.write(fun["code"])
             
 def extract_functions_from_file(input_file):