#include "maze.h"
#include <stdlib.h>
#include <string.h>

void maze_init(Maze *m, int rows, int cols) {
    m->rows = rows;
    m->cols = cols;
    maze_reset(m);
}

void maze_reset(Maze *m) {
    for (int r = 0; r < m->rows; r++)
        for (int c = 0; c < m->cols; c++) {
            m->cells[r][c] = W_N | W_E | W_S | W_W;
            m->visited[r][c] = 0;
        }
}

static void remove_wall(Maze *m, int r1, int c1, int r2, int c2) {
    int dr = r2 - r1, dc = c2 - c1;
    if (dr == -1) { m->cells[r1][c1] &= ~W_N; m->cells[r2][c2] &= ~W_S; }
    if (dr ==  1) { m->cells[r1][c1] &= ~W_S; m->cells[r2][c2] &= ~W_N; }
    if (dc == -1) { m->cells[r1][c1] &= ~W_W; m->cells[r2][c2] &= ~W_E; }
    if (dc ==  1) { m->cells[r1][c1] &= ~W_E; m->cells[r2][c2] &= ~W_W; }
}

static const int dx4[] = {0, 1, 0, -1};
static const int dy4[] = {-1, 0, 1, 0};

/* Fisher-Yates shuffle */
static void shuffle(int *arr, int n) {
    for (int i = n - 1; i > 0; i--) {
        int j = rand() % (i + 1);
        int t = arr[i]; arr[i] = arr[j]; arr[j] = t;
    }
}

/* Recursive backtracking - instant */
static void gen_recursive(Maze *m) {
    int stack[MAZE_MAX_ROWS * MAZE_MAX_COLS * 2];
    int sp = 0;
    stack[sp++] = 0; stack[sp++] = 0;
    m->visited[0][0] = 1;
    while (sp > 0) {
        int cy = stack[sp - 2], cx = stack[sp - 1];
        int dirs[] = {0, 1, 2, 3};
        shuffle(dirs, 4);
        int found = 0;
        for (int i = 0; i < 4; i++) {
            int ny = cy + dy4[dirs[i]], nx = cx + dx4[dirs[i]];
            if (ny >= 0 && ny < m->rows && nx >= 0 && nx < m->cols && !m->visited[ny][nx]) {
                m->visited[ny][nx] = 1;
                remove_wall(m, cy, cx, ny, nx);
                stack[sp++] = ny; stack[sp++] = nx;
                found = 1;
                break;
            }
        }
        if (!found) { sp -= 2; }
    }
}

/* Prim's algorithm */
static void gen_prim(Maze *m) {
    int sx = 0, sy = 0;
    m->visited[sy][sx] = 1;
    typedef struct { int x, y; } Edge;
    Edge frontier[MAZE_MAX_ROWS * MAZE_MAX_COLS * 4];
    int fc = 0;

    for (int d = 0; d < 4; d++) {
        int ny = sy + dy4[d], nx = sx + dx4[d];
        if (ny >= 0 && ny < m->rows && nx >= 0 && nx < m->cols)
            frontier[fc++] = (Edge){nx, ny};
    }

    while (fc > 0) {
        int idx = rand() % fc;
        Edge e = frontier[idx];
        frontier[idx] = frontier[--fc];
        if (m->visited[e.y][e.x]) continue;
        m->visited[e.y][e.x] = 1;

        int dirs[] = {0, 1, 2, 3};
        shuffle(dirs, 4);
        for (int i = 0; i < 4; i++) {
            int ny = e.y + dy4[dirs[i]], nx = e.x + dx4[dirs[i]];
            if (ny >= 0 && ny < m->rows && nx >= 0 && nx < m->cols && m->visited[ny][nx]) {
                remove_wall(m, e.y, e.x, ny, nx);
                break;
            }
        }

        for (int d = 0; d < 4; d++) {
            int ny = e.y + dy4[d], nx = e.x + dx4[d];
            if (ny >= 0 && ny < m->rows && nx >= 0 && nx < m->cols && !m->visited[ny][nx])
                frontier[fc++] = (Edge){nx, ny};
        }
    }
}

/* Kruskal's with Union-Find */
static int uf_parent[MAZE_MAX_ROWS * MAZE_MAX_COLS];
static int uf_rank[MAZE_MAX_ROWS * MAZE_MAX_COLS];

static int uf_find(int x) {
    while (uf_parent[x] != x) { uf_parent[x] = uf_parent[uf_parent[x]]; x = uf_parent[x]; }
    return x;
}

static int uf_union(int a, int b) {
    a = uf_find(a); b = uf_find(b);
    if (a == b) return 0;
    if (uf_rank[a] < uf_rank[b]) { int t = a; a = b; b = t; }
    uf_parent[b] = a;
    if (uf_rank[a] == uf_rank[b]) uf_rank[a]++;
    return 1;
}

static void gen_kruskal(Maze *m) {
    int total = m->rows * m->cols;
    for (int i = 0; i < total; i++) { uf_parent[i] = i; uf_rank[i] = 0; }

    typedef struct { int x1, y1, x2, y2; } Wall;
    Wall walls[MAZE_MAX_ROWS * MAZE_MAX_COLS * 2];
    int wc = 0;

    for (int r = 0; r < m->rows; r++)
        for (int c = 0; c < m->cols; c++) {
            if (r + 1 < m->rows) walls[wc++] = (Wall){c, r, c, r + 1};
            if (c + 1 < m->cols) walls[wc++] = (Wall){c, r, c + 1, r};
        }

    shuffle((int*)walls, wc);

    for (int i = 0; i < wc; i++) {
        int a = walls[i].y1 * m->cols + walls[i].x1;
        int b = walls[i].y2 * m->cols + walls[i].x2;
        if (uf_union(a, b)) {
            remove_wall(m, walls[i].y1, walls[i].x1, walls[i].y2, walls[i].x2);
            m->visited[walls[i].y1][walls[i].x1] = 1;
            m->visited[walls[i].y2][walls[i].x2] = 1;
        }
    }
}

/* Aldous-Broder (random walk) */
static void gen_aldous(Maze *m) {
    int cx = 0, cy = 0;
    int remaining = m->rows * m->cols - 1;
    m->visited[cy][cx] = 1;
    while (remaining > 0) {
        int d = rand() % 4;
        int nx = cx + dx4[d], ny = cy + dy4[d];
        if (nx < 0 || nx >= m->cols || ny < 0 || ny >= m->rows) continue;
        if (!m->visited[ny][nx]) {
            remove_wall(m, cy, cx, ny, nx);
            m->visited[ny][nx] = 1;
            remaining--;
        }
        cx = nx; cy = ny;
    }
}

int maze_generate_full(Maze *m, MazeAlgo algo) {
    maze_reset(m);
    switch (algo) {
    case ALG_RECURSIVE: gen_recursive(m); break;
    case ALG_PRIM:      gen_prim(m); break;
    case ALG_KRUSKAL:   gen_kruskal(m); break;
    case ALG_ALDUS:     gen_aldous(m); break;
    default:            gen_recursive(m); break;
    }
    return 1;
}
