#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/inotify.h>
#include <sys/wait.h>
#include <sys/stat.h>
#include <sys/file.h>
#include <signal.h>
#include <dirent.h>
#include <time.h>
#include <stdarg.h>
#include <errno.h>
#include <libgen.h>
#include <fcntl.h>
#include <limits.h>
#include <sys/prctl.h>

#define EVENT_SIZE          (sizeof(struct inotify_event))
#define EVENT_BUF_LEN       (1024 * (EVENT_SIZE + 16))
#define LOG_MAX_SIZE        (10 * 1024 * 1024)
#define CRASH_RESTART_DELAY 15
#define TERM_TIMEOUT        5
#define APP_LOG_FILE        "rproc.log"
#define PID_LOCK_FILE       "/tmp/rproc.pid"
#define MAX_LOG_LINE_LEN    4096
#define TAIL_LINE_COUNT     300

typedef struct Process {
    pid_t pid;
    char* script_name;
    time_t crashed_at;
    struct Process* next;
} Process;

static Process* process_list = NULL;
static volatile sig_atomic_t running = 1;

static void app_log(const char* format, ...);
static void terminate_process(pid_t pid, const char* script_name);
static void remove_process_by_pid(pid_t pid);
static void start_script(const char* script_name);

static void check_and_truncate_file(const char* path, off_t max_size) {
    struct stat st;
    if (stat(path, &st) == 0) {
        if (st.st_size > max_size) {
            if (truncate(path, 0) == -1) {
                fprintf(stderr, "rproc warning: could not truncate %s: %s\n", path, strerror(errno));
            }
        }
    }
}

static void app_log(const char* format, ...) {
    check_and_truncate_file(APP_LOG_FILE, LOG_MAX_SIZE);
    FILE* fp = fopen(APP_LOG_FILE, "a");
    if (!fp) return;

    time_t now = time(NULL);
    char time_buf[20];
    strftime(time_buf, sizeof(time_buf), "%Y-%m-%d %H:%M:%S", localtime(&now));

    fprintf(fp, "[%s] ", time_buf);
    va_list args;
    va_start(args, format);
    vfprintf(fp, format, args);
    va_end(args);
    fputc('\n', fp);
    fclose(fp);
}

static void add_process(pid_t pid, const char* script_name) {
    Process* p = malloc(sizeof(Process));
    if (!p) {
        app_log("FATAL: malloc failed in add_process.");
        exit(EXIT_FAILURE);
    }
    p->pid = pid;
    p->script_name = strdup(script_name);
    if (!p->script_name) {
        app_log("FATAL: strdup failed for script %s.", script_name);
        free(p);
        exit(EXIT_FAILURE);
    }
    p->crashed_at = 0;
    p->next = process_list;
    process_list = p;
}

static Process* find_process_by_pid(pid_t pid) {
    for (Process* p = process_list; p; p = p->next) {
        if (p->pid == pid) return p;
    }
    return NULL;
}

static Process* find_process_by_name(const char* script_name) {
    for (Process* p = process_list; p; p = p->next) {
        if (strcmp(p->script_name, script_name) == 0) return p;
    }
    return NULL;
}

static void remove_process_by_pid(pid_t pid) {
    Process** current = &process_list;
    while (*current) {
        Process* entry = *current;
        if (entry->pid == pid) {
            *current = entry->next;
            free(entry->script_name);
            free(entry);
            return;
        }
        current = &entry->next;
    }
}

static void terminate_process(pid_t pid, const char* script_name) {
    if (pid <= 0) {
        if (pid == -1) {
            remove_process_by_pid(pid);
        }
        return;
    }

    app_log("Stopping process for '%s' (PID: %d)", script_name, pid);
    if (kill(pid, SIGTERM) == -1) {
        if (errno == ESRCH) {
            app_log("Process %d for '%s' did not exist. Removing from tracking.", pid, script_name);
            char pid_filename[FILENAME_MAX];
            snprintf(pid_filename, sizeof(pid_filename), "%s.pid", script_name);
            if (access(pid_filename, F_OK) == 0) {
                unlink(pid_filename);
            }
            remove_process_by_pid(pid);
        }
        return;
    }

    for (int i = 0; i < TERM_TIMEOUT; ++i) {
        int status;
        const pid_t result = waitpid(pid, &status, WNOHANG);
        if (result == pid) {
            app_log("Process %d for '%s' terminated gracefully.", pid, script_name);
            char pid_filename[FILENAME_MAX];
            snprintf(pid_filename, sizeof(pid_filename), "%s.pid", script_name);
            if (access(pid_filename, F_OK) == 0) {
                unlink(pid_filename);
            }
            remove_process_by_pid(pid);
            return;
        }
        if (result == -1 && errno != ECHILD) {
            remove_process_by_pid(pid);
            return;
        }
        sleep(1);
    }

    app_log("Process %d for '%s' did not terminate in time. Sending SIGKILL.", pid, script_name);
    if (kill(pid, SIGKILL) == -1 && errno == ESRCH) {
        app_log("Process %d for '%s' disappeared before SIGKILL.", pid, script_name);
    }
    waitpid(pid, NULL, 0);
    char pid_filename[FILENAME_MAX];
    snprintf(pid_filename, sizeof(pid_filename), "%s.pid", script_name);
    if (access(pid_filename, F_OK) == 0) {
        unlink(pid_filename);
    }
    remove_process_by_pid(pid);
}

static void create_slug_from_script_name(const char* script_name, char* slug_buffer, size_t buffer_size) {
    strncpy(slug_buffer, script_name, buffer_size - 1);
    slug_buffer[buffer_size - 1] = '\0';

    char* dot = strrchr(slug_buffer, '.');
    if (dot && strcmp(dot, ".sh") == 0) {
        *dot = '\0';
    }

    for (char* p = slug_buffer; *p; ++p) {
        if (*p == '_') {
            *p = '-';
        }
    }
}

static void start_script(const char* script_name) {
    Process *existing = find_process_by_name(script_name);
    if (existing && existing->pid > 0) {
        app_log("Script '%s' is already running with PID %d.", script_name, existing->pid);
        return;
    }

    app_log("Starting script: '%s'", script_name);
    const pid_t pid = fork();
    if (pid == -1) {
        app_log("Failed to fork for script '%s': %s", script_name, strerror(errno));
        return;
    }

    if (pid == 0) {
        char slug[128];
        char new_proc_name[256];
        create_slug_from_script_name(script_name, slug, sizeof(slug));
        snprintf(new_proc_name, sizeof(new_proc_name), "rproc-%s-process", slug);
        
        if (prctl(PR_SET_NAME, new_proc_name, 0, 0, 0) == -1) {
            perror("prctl");
        }

        char stdout_log[FILENAME_MAX], stderr_log[FILENAME_MAX];
        snprintf(stdout_log, sizeof(stdout_log), "%s.stdout.log", script_name);
        snprintf(stderr_log, sizeof(stderr_log), "%s.stderr.log", script_name);

        check_and_truncate_file(stdout_log, LOG_MAX_SIZE);
        check_and_truncate_file(stderr_log, LOG_MAX_SIZE);
        
        if (!freopen(stdout_log, "a", stdout) || !freopen(stderr_log, "a", stderr)) {
             exit(127);
        }

        char* const argv[] = {"/bin/bash", (char*)script_name, NULL};
        char* const envp[] = { "PATH=/usr/bin:/bin", NULL };
        execve("/bin/bash", argv, envp);
        
        fprintf(stderr, "Failed to execute script '%s': %s\n", script_name, strerror(errno));
        exit(127);
    }
    
    char pid_filename[FILENAME_MAX];
    snprintf(pid_filename, sizeof(pid_filename), "%s.pid", script_name);
    FILE* pid_fp = fopen(pid_filename, "w");
    if (pid_fp) {
        fprintf(pid_fp, "%d", pid);
        fclose(pid_fp);
    } else {
        app_log("Warning: could not create pid file for %s", script_name);
    }

    add_process(pid, script_name);
    app_log("Started '%s' with PID %d", script_name, pid);
}

static void stop_script(const char* script_name) {
    Process* p = find_process_by_name(script_name);
    if (p) {
        terminate_process(p->pid, p->script_name);
    }
}

static void handle_sigchld(int sig) {
    (void)sig;
    int saved_errno = errno;
    int status;
    pid_t pid;

    while ((pid = waitpid(-1, &status, WNOHANG)) > 0) {
        Process* p = find_process_by_pid(pid);
        if (p) {
            char pid_filename[FILENAME_MAX];
            snprintf(pid_filename, sizeof(pid_filename), "%s.pid", p->script_name);
            
            if (WIFEXITED(status) && WEXITSTATUS(status) != 0) {
                app_log("Script '%s' (PID: %d) crashed with status %d.", p->script_name, pid, WEXITSTATUS(status));
                if (access(pid_filename, F_OK) == 0) {
                    unlink(pid_filename);
                }
                p->crashed_at = time(NULL);
                p->pid = -1;
            } else if (WIFSIGNALED(status)) {
                app_log("Script '%s' (PID: %d) terminated by signal %d.", p->script_name, pid, WTERMSIG(status));
                if (access(pid_filename, F_OK) == 0) {
                    unlink(pid_filename);
                }
                p->crashed_at = time(NULL);
                p->pid = -1;
            } else {
                 app_log("Script '%s' (PID: %d) exited cleanly.", p->script_name, pid);
                 if (access(pid_filename, F_OK) == 0) {
                    unlink(pid_filename);
                 }
                 remove_process_by_pid(pid);
            }
        }
    }
    errno = saved_errno;
}

static void handle_shutdown_signals(int sig) {
    (void)sig;
    char msg[] = "Shutdown signal received.\n";
    ssize_t bytes_written = write(STDOUT_FILENO, msg, sizeof(msg) - 1);
    (void)bytes_written;
    running = 0;
}

static void setup_signal_handlers(void) {
    struct sigaction sa_chld = {0}, sa_term = {0};
    sa_chld.sa_handler = handle_sigchld;
    sa_chld.sa_flags = SA_RESTART | SA_NOCLDSTOP;
    sigemptyset(&sa_chld.sa_mask);
    if (sigaction(SIGCHLD, &sa_chld, 0) == -1) {
        perror("sigaction SIGCHLD");
        exit(EXIT_FAILURE);
    }

    sa_term.sa_handler = handle_shutdown_signals;
    sigemptyset(&sa_term.sa_mask);
    if (sigaction(SIGTERM, &sa_term, 0) == -1 || sigaction(SIGINT, &sa_term, 0) == -1) {
        perror("sigaction SIGTERM/SIGINT");
        exit(EXIT_FAILURE);
    }
}

static void adopt_existing_processes(void) {
    DIR* d = opendir(".");
    if (!d) return;

    struct dirent* dir;
    while ((dir = readdir(d)) != NULL) {
        size_t len = strlen(dir->d_name);
        if (len > 4 && strcmp(dir->d_name + len - 4, ".pid") == 0) {
            FILE* fp = fopen(dir->d_name, "r");
            if (!fp) continue;

            pid_t pid = 0;
            if (fscanf(fp, "%d", &pid) == 1 && pid > 0) {
                fclose(fp);
                if (kill(pid, 0) == 0 || errno == EPERM) {
                    char script_name[FILENAME_MAX];
                    strncpy(script_name, dir->d_name, len - 4);
                    script_name[len - 4] = '\0';
                    app_log("Adopting existing process '%s' with PID %d", script_name, pid);
                    add_process(pid, script_name);
                } else {
                    app_log("Found stale pid file '%s'. Removing.", dir->d_name);
                    if (access(dir->d_name, F_OK) == 0) {
                        unlink(dir->d_name);
                    }
                }
            } else {
                fclose(fp);
            }
        }
    }
    closedir(d);
}

static void run_daemon(void) {
    adopt_existing_processes();

    DIR* d = opendir(".");
    if (d) {
        struct dirent* dir;
        while ((dir = readdir(d)) != NULL) {
            size_t len = strlen(dir->d_name);
            if (len > 3 && strcmp(dir->d_name + len - 3, ".sh") == 0) {
                start_script(dir->d_name);
            }
        }
        closedir(d);
    }

    int fd = inotify_init1(IN_NONBLOCK);
    if (fd < 0) {
        app_log("inotify_init1 failed: %s", strerror(errno));
        return;
    }
    inotify_add_watch(fd, ".", IN_CREATE | IN_DELETE | IN_MODIFY | IN_MOVED_TO | IN_MOVED_FROM);
    app_log("Daemon started. Monitoring current directory.");

    while (running) {
        fd_set read_fds;
        FD_ZERO(&read_fds);
        FD_SET(fd, &read_fds);
        struct timeval tv = { .tv_sec = 1, .tv_usec = 0 };

        int retval = select(fd + 1, &read_fds, NULL, NULL, &tv);

        if (retval > 0 && FD_ISSET(fd, &read_fds)) {
            char buffer[EVENT_BUF_LEN];
            ssize_t length = read(fd, buffer, EVENT_BUF_LEN);
            if (length > 0) {
                for (int i = 0; i < length; ) {
                    struct inotify_event* event = (struct inotify_event*)&buffer[i];
                    if (event->len) {
                        size_t name_len = strlen(event->name);
                        if (name_len > 3 && strcmp(event->name + name_len - 3, ".sh") == 0) {
                            if (event->mask & (IN_CREATE | IN_MOVED_TO)) {
                                app_log("Detected new script: '%s'", event->name);
                                start_script(event->name);
                            } else if (event->mask & IN_MODIFY) {
                                app_log("Detected modification: '%s'. Restarting.", event->name);
                                stop_script(event->name);
                                start_script(event->name);
                            } else if (event->mask & (IN_DELETE | IN_MOVED_FROM)) {
                                app_log("Detected deletion: '%s'. Stopping.", event->name);
                                stop_script(event->name);
                            }
                        }
                    }
                    i += EVENT_SIZE + event->len;
                }
            }
        }

        Process* p = process_list;
        while (p) {
            Process* next = p->next;
            if (p->pid == -1 && p->crashed_at > 0) {
                if (time(NULL) - p->crashed_at >= CRASH_RESTART_DELAY) {
                    app_log("Restarting crashed script '%s'.", p->script_name);
                    char* name_copy = strdup(p->script_name);
                    if (name_copy) {
                        remove_process_by_pid(p->pid); 
                        start_script(name_copy);
                        free(name_copy);
                    }
                }
            }
            p = next;
        }
    }

    app_log("Shutdown sequence initiated. Terminating all processes.");
    while (process_list) {
        terminate_process(process_list->pid, process_list->script_name);
    }
    close(fd);
    app_log("Daemon stopped.");
}

typedef struct LogFile {
    char* name;
    off_t offset;
    int color_index;
    struct LogFile* next;
} LogFile;

static off_t get_tail_offset(const char* filename, int line_count) {
    FILE* fp = fopen(filename, "rb");
    if (!fp) return 0;

    char buf[4096];
    off_t file_size;
    long pos;
    int newlines = 0;

    if (fseek(fp, 0, SEEK_END) != 0) {
        fclose(fp);
        return 0;
    }
    file_size = ftell(fp);
    if (file_size <= 0) {
        fclose(fp);
        return 0;
    }
    pos = file_size;

    while (pos > 0) {
        long seek_to = pos - sizeof(buf);
        if (seek_to < 0) seek_to = 0;
        
        if (fseek(fp, seek_to, SEEK_SET) != 0) break;
        
        size_t bytes_read = fread(buf, 1, pos - seek_to, fp);
        if (bytes_read == 0) break;

        for (long i = bytes_read - 1; i >= 0; --i) {
            if (buf[i] == '\n' && (seek_to + i) != (file_size - 1)) {
                newlines++;
                if (newlines >= line_count) {
                    pos = seek_to + i + 1;
                    fclose(fp);
                    return pos;
                }
            }
        }
        pos = seek_to;
    }

    fclose(fp);
    return 0;
}


static void run_monitor(void) {
    char cwd[PATH_MAX] = {0};
    char line_buffer[PATH_MAX];

    FILE* lock_fp = fopen(PID_LOCK_FILE, "r");
    if (lock_fp) {
        if (fgets(line_buffer, sizeof(line_buffer), lock_fp)) {
            if (fgets(cwd, sizeof(cwd), lock_fp)) {
                cwd[strcspn(cwd, "\n")] = 0;
                if (chdir(cwd) == -1) {
                    perror("Could not change to daemon's directory");
                    fclose(lock_fp);
                    return;
                }
            }
        }
        fclose(lock_fp);
    }

    printf("Another instance is running in %s. Attaching as a live monitor...\n", cwd);
    printf("--- Tailing last %d lines from all *.log files. Press Ctrl+C to exit. ---\n", TAIL_LINE_COUNT);

    LogFile* log_list = NULL;
    int fd = inotify_init();
    if (fd < 0) { perror("inotify_init"); return; }
    
    inotify_add_watch(fd, ".", IN_CREATE | IN_MODIFY);
    
    static const char* colors[] = {
        "\033[0;32m", /* Green */
        "\033[0;33m", /* Yellow */
        "\033[0;34m", /* Blue */
        "\033[0;36m", /* Cyan */
        "\033[0;35m", /* Magenta */
        "\033[0;31m", /* Red */
    };
    static const int num_colors = sizeof(colors) / sizeof(colors[0]);
    static const char* color_reset = "\033[0m";
    static int next_color_index = 0;


    while(1) {
        DIR* d = opendir(".");
        if (d) {
            struct dirent* dir;
            while ((dir = readdir(d)) != NULL) {
                size_t len = strlen(dir->d_name);
                if (len > 4 && strcmp(dir->d_name + len - 4, ".log") == 0) {
                    int found = 0;
                    for (LogFile* l = log_list; l; l = l->next) {
                        if (strcmp(l->name, dir->d_name) == 0) {
                            found = 1;
                            break;
                        }
                    }
                    if (!found) {
                        LogFile* new_log = malloc(sizeof(LogFile));
                        new_log->name = strdup(dir->d_name);
                        new_log->offset = get_tail_offset(dir->d_name, TAIL_LINE_COUNT);
                        new_log->color_index = next_color_index;
                        next_color_index = (next_color_index + 1) % num_colors;
                        new_log->next = log_list;
                        log_list = new_log;
                    }
                }
            }
            closedir(d);
        }

        for (LogFile* l = log_list; l; l = l->next) {
             FILE* fp = fopen(l->name, "r");
             if (fp) {
                 fseek(fp, l->offset, SEEK_SET);
                 char line[MAX_LOG_LINE_LEN];
                 while (fgets(line, sizeof(line), fp)) {
                     line[strcspn(line, "\r\n")] = 0;
                     printf("%s[%s] %s%s\n", colors[l->color_index], l->name, line, color_reset);
                 }
                 l->offset = ftell(fp);
                 fclose(fp);
             }
        }
        
        fd_set read_fds;
        FD_ZERO(&read_fds);
        FD_SET(fd, &read_fds);
        struct timeval tv = { .tv_sec = 2, .tv_usec = 0 };
        select(fd + 1, &read_fds, NULL, NULL, &tv);

        if (FD_ISSET(fd, &read_fds)) {
            char buffer[EVENT_BUF_LEN];
            ssize_t bytes_read = read(fd, buffer, EVENT_BUF_LEN);
            (void)bytes_read;
        }
    }
}

int main(void) {
    int pid_fd = open(PID_LOCK_FILE, O_CREAT | O_RDWR, 0666);
    if (pid_fd == -1) {
        perror("Could not open PID lock file");
        return 1;
    }

    if (flock(pid_fd, LOCK_EX | LOCK_NB) == -1) {
        if (errno == EWOULDBLOCK) {
            close(pid_fd);
            run_monitor();
        } else {
            perror("flock");
            close(pid_fd);
            return 1;
        }
    } else {
        if (ftruncate(pid_fd, 0) == -1) {
            perror("Could not truncate PID lock file");
            goto cleanup_lock;
        }

        char cwd[PATH_MAX];
        if (!getcwd(cwd, sizeof(cwd))) {
            perror("Could not get current working directory");
            goto cleanup_lock;
        }

        char lock_content[PATH_MAX + 16];
        int len = snprintf(lock_content, sizeof(lock_content), "%d\n%s\n", getpid(), cwd);
        if (write(pid_fd, lock_content, len) != (long int)len) {
            perror("Could not write to PID lock file");
            goto cleanup_lock;
        }

        setup_signal_handlers();
        run_daemon();

    cleanup_lock:
        flock(pid_fd, LOCK_UN);
        close(pid_fd);
        if (access(PID_LOCK_FILE, F_OK) == 0) {
            unlink(PID_LOCK_FILE);
        }
    }

    return 0;
}