2 files changed, 15 insertions, 5 deletions

diff --git a/list.c b/list.c
index 8c5b478..e1d4ffd 100644
--- a/list.c
+++ b/list.c
@@ -23,6 +23,17 @@ void add_file_path(Node **head, char *file_path) {
 }
 
 void list_files_recursively(char *base_path, Node **head) {
+        struct stat statbuf;
+        if (stat(base_path, &statbuf) == -1) {
+                perror("stat");
+                return;
+        }
+
+        if (S_ISREG(statbuf.st_mode)) {
+                add_file_path(head, base_path);
+                return;
+        }
+
         char path[2048];
         struct dirent *dp;
         DIR *dir = opendir(base_path);
@@ -39,11 +50,10 @@ void list_files_recursively(char *base_path, Node **head) {
                                 continue;
                         }
 
-                        struct stat statbuf;
                         if (stat(path, &statbuf) != -1) {
                                 if (S_ISDIR(statbuf.st_mode)) {
                                         list_files_recursively(path, head);
-                                } else {
+                                } else if (S_ISREG(statbuf.st_mode)) {
                                         add_file_path(head, path);
                                 }
                         }
diff --git a/main.c b/main.c
index 0a2ddd8..4ff4ab5 100644
--- a/main.c
+++ b/main.c
@@ -177,13 +177,13 @@ const char *get_file_extension(const char *file_path) {
 }
 
 int main(int argc, char *argv[]) {
-        if (argc < 3) {
-                printf("Usage: %s <search term> <directory>\n", argv[0]);
+        if (argc < 2) {
+                printf("Usage: %s <search term> [directory|file]\n", argv[0]);
                 return 1;
         }
 
         const char *cfname = argv[1];
-        char *directory = argv[2];
+        char *directory = (argc > 2) ? argv[2] : ".";
 
         TSLanguage *tree_sitter_c(void);
         TSLanguage *tree_sitter_python(void);