// retoor <retoor@molodetz.nl>

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

#define INDENT_WIDTH 4

static char *dedent_code(const char *src) {
    if (!src || !*src) return strdup("");

    int min_indent = -1;
    const char *line = src;
    while (line && *line) {
        int indent = 0;
        while (line[indent] == ' ' || line[indent] == '\t') indent++;
        
        if (line[indent] != '\n' && line[indent] != '\0') {
            if (min_indent == -1 || indent < min_indent) min_indent = indent;
        }
        
        line = strchr(line, '\n');
        if (line) line++;
    }

    if (min_indent <= 0) return strdup(src);

    size_t src_len = strlen(src);
    char *result = malloc(src_len + 1);
    if (!result) return strdup(src);
    
    char *dst = result;
    const char *curr = src;
    while (curr && *curr) {
        int to_skip = min_indent;
        while (to_skip > 0 && (*curr == ' ' || *curr == '\t')) {
            curr++;
            to_skip--;
        }
        
        const char *next_line = strchr(curr, '\n');
        if (next_line) {
            size_t line_len = (size_t)(next_line - curr + 1);
            memcpy(dst, curr, line_len);
            dst += line_len;
            curr = next_line + 1;
        } else {
            strcpy(dst, curr);
            break;
        }
    }
    *dst = '\0';
    return result;
}

static char *normalize_indentation(const char *src) {
    if (!src) return NULL;
    size_t src_len = strlen(src);
    char *result = malloc(src_len * 2 + 1); // Extra space for normalized indents
    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);

        // Check if line is empty or just whitespace
        bool is_empty = true;
        for (size_t i = 0; i < line_len; i++) {
            if (!isspace((unsigned char)line[i])) {
                is_empty = false;
                break;
            }
        }

        if (is_empty) {
            if (next_line) {
                *dst++ = '\n';
                line = next_line + 1;
            } else {
                break;
            }
            continue;
        }

        // Calculate current leading indent
        int leading_spaces = 0;
        const char *content_ptr = line;
        while (content_ptr < (line + line_len) && (*content_ptr == ' ' || *content_ptr == '\t')) {
            if (*content_ptr == '\t') {
                leading_spaces += INDENT_WIDTH;
            } else {
                leading_spaces++;
            }
            content_ptr++;
        }

        // Round to nearest INDENT_WIDTH
        int normalized_level = (leading_spaces + INDENT_WIDTH / 2) / INDENT_WIDTH;
        int target_spaces = normalized_level * INDENT_WIDTH;

        for (int i = 0; i < target_spaces; i++) *dst++ = ' ';
        
        size_t content_len = (size_t)((line + line_len) - content_ptr);
        memcpy(dst, content_ptr, content_len);
        dst += content_len;

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

static char *repair_strings_and_brackets(const char *src) {
    if (!src) return NULL;
    size_t src_len = strlen(src);
    char *result = malloc(src_len + 2048); // Buffer for extra quotes/brackets
    if (!result) return NULL;

    char *dst = result;
    const char *p = src;
    
    char bracket_stack[1024];
    int stack_ptr = 0;

    bool in_string = false;
    char string_quote = 0;
    int quote_type = 0; // 1 for single, 3 for triple
    bool escaped = false;

    while (*p) {
        char ch = *p;

        if (!in_string) {
            if (ch == '#') {
                // Comment, copy until newline
                while (*p && *p != '\n') *dst++ = *p++;
                continue;
            }
            if (ch == '\'' || ch == '"') {
                string_quote = ch;
                if (strncmp(p, "'''", 3) == 0 || strncmp(p, "\"\"\"", 3) == 0) {
                    quote_type = 3;
                    in_string = true;
                    *dst++ = *p++; *dst++ = *p++; *dst++ = *p++;
                    continue;
                } else {
                    quote_type = 1;
                    in_string = true;
                    *dst++ = *p++;
                    continue;
                }
            } else if (ch == '(' || ch == '[' || ch == '{') {
                if (stack_ptr < 1024) bracket_stack[stack_ptr++] = ch;
            } else if (ch == ')' || ch == ']' || ch == '}') {
                char expected = 0;
                if (ch == ')') expected = '(';
                else if (ch == ']') expected = '[';
                else if (ch == '}') expected = '{';
                
                if (stack_ptr > 0 && bracket_stack[stack_ptr - 1] == expected) {
                    stack_ptr--;
                } else {
                    // Mismatched closing; skip it to prevent syntax errors
                    p++;
                    continue;
                }
            }
        } else {
            if (escaped) {
                escaped = false;
            } else if (ch == '\\') {
                escaped = true;
            } else if (ch == string_quote) {
                if (quote_type == 3) {
                    if (strncmp(p, "'''", 3) == 0 && string_quote == '\'') {
                        in_string = false;
                        *dst++ = *p++; *dst++ = *p++; *dst++ = *p++;
                        continue;
                    } else if (strncmp(p, "\"\"\"", 3) == 0 && string_quote == '"') {
                        in_string = false;
                        *dst++ = *p++; *dst++ = *p++; *dst++ = *p++;
                        continue;
                    }
                } else {
                    in_string = false;
                }
            }
        }
        *dst++ = *p++;
    }

    if (in_string) {
        if (quote_type == 3) {
            *dst++ = string_quote; *dst++ = string_quote; *dst++ = string_quote;
        } else {
            *dst++ = string_quote;
        }
    }

    // Balance brackets
    while (stack_ptr > 0) {
        char opener = bracket_stack[--stack_ptr];
        if (opener == '(') *dst++ = ')';
        else if (opener == '[') *dst++ = ']';
        else if (opener == '{') *dst++ = '}';
    }

    *dst = '\0';
    return result;
}

static char *add_missing_passes(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 *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);

        // Copy current line
        memcpy(dst, line, line_len);
        dst += line_len;
        if (next_line) *dst++ = '\n';

        // Check if line ends with ':' (ignoring comments/whitespace)
        const char *p = line + line_len - 1;
        while (p >= line && isspace((unsigned char)*p)) p--;
        
        if (p >= line && *p == ':') {
            // Heuristic: check if next non-empty line is more indented
            const char *lookahead = next_line ? next_line + 1 : NULL;
            bool needs_pass = true;
            
            while (lookahead && *lookahead) {
                // Skip empty/whitespace lines
                const char *next_next = strchr(lookahead, '\n');
                size_t look_len = next_next ? (size_t)(next_next - lookahead) : strlen(lookahead);
                
                bool empty = true;
                for (size_t i = 0; i < look_len; i++) {
                    if (!isspace((unsigned char)lookahead[i])) {
                        empty = false;
                        break;
                    }
                }
                
                if (!empty) {
                    // Check indent of this non-empty line
                    int current_indent = 0;
                    const char *line_p = line;
                    while (line_p < (line + line_len) && (*line_p == ' ' || *line_p == '\t')) {
                        if (*line_p == '\t') current_indent += INDENT_WIDTH;
                        else current_indent++;
                        line_p++;
                    }

                    int line_look_indent = 0;
                    const char *look_p = lookahead;
                    while (look_p < (lookahead + look_len) && (*look_p == ' ' || *look_p == '\t')) {
                        if (*look_p == '\t') line_look_indent += INDENT_WIDTH;
                        else line_look_indent++;
                        look_p++;
                    }

                    if (line_look_indent > current_indent) {
                        needs_pass = false;
                    }
                    break;
                }
                
                if (next_next) lookahead = next_next + 1;
                else break;
            }

            if (needs_pass) {
                // Find current indent to place 'pass' correctly
                int current_indent = 0;
                const char *line_p = line;
                while (line_p < (line + line_len) && (*line_p == ' ' || *line_p == '\t')) {
                    if (*line_p == '\t') current_indent += INDENT_WIDTH;
                    else current_indent++;
                    line_p++;
                }
                
                int target_indent = current_indent + INDENT_WIDTH;
                for (int i = 0; i < target_indent; i++) *dst++ = ' ';
                memcpy(dst, "pass\n", 5);
                dst += 5;
            }
        }

        if (next_line) line = next_line + 1;
        else break;
    }

    *dst = '\0';
    return result;
}

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

    char *s1 = dedent_code(src);
    char *s2 = normalize_indentation(s1);
    free(s1);
    
    char *s3 = repair_strings_and_brackets(s2);
    free(s2);
    
    char *s4 = add_missing_passes(s3);
    free(s3);

    return s4;
}