// retoor <retoor@molodetz.nl>

#include "bash_repair.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>
#include <ctype.h>

static char *ensure_shebang(const char *text) {
    if (!text || !*text) return strdup("");
    
    // Check if it already has a shebang
    const char *p = text;
    while (*p && isspace((unsigned char)*p)) p++;
    if (strncmp(p, "#!", 2) == 0 || strncmp(p, ": <<", 4) == 0) {
        return strdup(text);
    }

    // Heuristic: if it has multiple lines, add shebang
    if (strchr(text, '\n')) {
        char *result = malloc(strlen(text) + 32);
        if (!result) return strdup(text);
        strcpy(result, "#!/usr/bin/env bash\n");
        strcat(result, text);
        return result;
    }

    return strdup(text);
}

static char *normalize_whitespace_and_operators(const char *src) {
    if (!src) return NULL;
    size_t src_len = strlen(src);
    char *result = malloc(src_len * 2 + 1);
    if (!result) return NULL;

    char *dst = result;
    const char *curr = src;

    while (*curr) {
        if (*curr == '\r') {
            curr++;
            continue;
        }

        // Detect operators to normalize spaces around them
        const char *ops[] = {"||", "&&", ">>", "|&", "|", ";", ">", "<", NULL};
        bool matched_op = false;
        for (int i = 0; ops[i]; i++) {
            size_t op_len = strlen(ops[i]);
            if (strncmp(curr, ops[i], op_len) == 0) {
                // Remove preceding spaces in dst if any
                while (dst > result && isspace((unsigned char)*(dst - 1)) && *(dst - 1) != '\n') dst--;
                
                if (dst > result && *(dst - 1) != '\n') *dst++ = ' ';
                memcpy(dst, ops[i], op_len);
                dst += op_len;
                curr += op_len;
                
                // Skip following spaces in curr
                while (*curr && isspace((unsigned char)*curr) && *curr != '\n') curr++;
                if (*curr && *curr != '\n') *dst++ = ' ';
                
                matched_op = true;
                break;
            }
        }

        if (!matched_op) {
            *dst++ = *curr++;
        }
    }
    *dst = '\0';

    // Second pass to strip trailing spaces on each line
    char *s2 = strdup(result);
    free(result);
    if (!s2) return NULL;

    char *final = malloc(strlen(s2) + 1);
    char *f_ptr = final;
    char *line = s2;
    while (line && *line) {
        char *next_line = strchr(line, '\n');
        size_t line_len = next_line ? (size_t)(next_line - line) : strlen(line);
        
        const char *end = line + line_len - 1;
        while (end >= line && isspace((unsigned char)*end)) end--;
        
        size_t new_len = (size_t)(end - line + 1);
        memcpy(f_ptr, line, new_len);
        f_ptr += new_len;
        if (next_line) {
            *f_ptr++ = '\n';
            line = next_line + 1;
        } else {
            break;
        }
    }
    *f_ptr = '\0';
    free(s2);
    return final;
}

static char *fix_line_issues(const char *src) {
    if (!src) return NULL;
    size_t src_len = strlen(src);
    char *result = malloc(src_len * 2 + 1024);
    if (!result) return NULL;

    char *dst = result;
    const char *line = src;

    while (line && *line) {
        const char *next_line = strchr(line, '\n');
        size_t line_len = next_line ? (size_t)(next_line - line) : strlen(line);

        char line_buf[4096];
        if (line_len >= sizeof(line_buf)) line_len = sizeof(line_buf) - 1;
        memcpy(line_buf, line, line_len);
        line_buf[line_len] = '\0';

        // 1. Tiny Shell Lint (else if -> elif)
        char *else_if = strstr(line_buf, "else if ");
        if (else_if) {
            // Very basic replacement
            memmove(else_if + 4, else_if + 8, strlen(else_if + 8) + 1);
            memcpy(else_if, "elif", 4);
        }

        // 2. Fix unbalanced quotes
        int single = 0, double_q = 0;
        bool escaped = false;
        for (size_t i = 0; i < strlen(line_buf); i++) {
            if (escaped) { escaped = false; continue; }
            if (line_buf[i] == '\\') escaped = true;
            else if (line_buf[i] == '\'') single++;
            else if (line_buf[i] == '"') double_q++;
        }
        if (single % 2 == 1 && double_q == 0) strcat(line_buf, "'");
        else if (double_q % 2 == 1 && single == 0) strcat(line_buf, "\"");

        // 3. Fix trailing operators
        size_t cur_len = strlen(line_buf);
        const char *ops[] = {"||", "&&", ">>", "|&", "|", ">", "<", NULL};
        for (int i = 0; ops[i]; i++) {
            size_t op_len = strlen(ops[i]);
            if (cur_len >= op_len) {
                if (strcmp(line_buf + cur_len - op_len, ops[i]) == 0) {
                    // Comment it
                    char temp[4096];
                    strcpy(temp, line_buf);
                    temp[cur_len - op_len] = '\0';
                    strcat(temp, "# ");
                    strcat(temp, ops[i]);
                    strcpy(line_buf, temp);
                    break;
                }
            }
        }

        // 4. Dangerous rm -rf check
        if (strstr(line_buf, "sudo rm -rf /") || strstr(line_buf, "rm -rf / ")) {
            strcpy(dst, "# WARNING: potentially destructive command detected\n");
            dst += strlen(dst);
        }

        strcpy(dst, line_buf);
        dst += strlen(dst);
        if (next_line) *dst++ = '\n';

        if (next_line) line = next_line + 1;
        else break;
    }
    *dst = '\0';
    return result;
}

static char *collapse_nested_bash_c(const char *src) {
    if (!src) return NULL;
    // Pattern: bash -c "bash -c '...'")
    // We'll just do a very basic string replacement for common variants
    char *s1 = strdup(src);
    char *patterns[] = {
        "bash -c \"bash -c '",
        "bash -c 'bash -c \"",
        NULL
    };
    
    for (int i = 0; patterns[i]; i++) {
        char *p;
        while ((p = strstr(s1, patterns[i]))) {
            // Find closing quotes
            char outer = patterns[i][8]; // " or '
            char inner = patterns[i][18]; // ' or "
            
            char *inner_end = strchr(p + 19, inner);
            if (inner_end && *(inner_end + 1) == outer) {
                // We can collapse. 
                // Original: [p]bash -c "bash -c 'cmd'"[end]
                // New:      [p]bash -c 'cmd'[end]
                size_t cmd_len = (size_t)(inner_end - (p + 19));
                char *new_s = malloc(strlen(s1) + 1);
                size_t prefix_len = (size_t)(p - s1);
                memcpy(new_s, s1, prefix_len);
                char *d = new_s + prefix_len;
                strcpy(d, "bash -c ");
                d += 8;
                *d++ = inner;
                memcpy(d, p + 19, cmd_len);
                d += cmd_len;
                *d++ = inner;
                strcpy(d, inner_end + 2);
                
                free(s1);
                s1 = new_s;
            } else {
                break; 
            }
        }
    }
    return s1;
}

char *bash_repair_command(const char *src) {
    if (!src) return NULL;

    char *s1 = normalize_whitespace_and_operators(src);
    char *s2 = fix_line_issues(s1);
    free(s1);
    char *s3 = collapse_nested_bash_c(s2);
    free(s2);
    char *s4 = ensure_shebang(s3);
    free(s3);

    return s4;
}