#include "tui.h"
#include <ncurses.h>
#include <string.h>
#include <stdlib.h>
#include <signal.h>

#define GRID_ROWS 8
#define GRID_COLS 8
#define CELL_WIDTH 6
#define CELL_HEIGHT 3

// Color pairs
enum {
    COLOR_EMPTY = 1,
    COLOR_RECORDING,
    COLOR_LOOPING,
    COLOR_STOPPED,
    COLOR_SELECTED,
    COLOR_HELP
};

static Engine *g_engine = NULL;
static int selected_row = 0;
static int selected_col = 0;
static bool show_help = false;

// Modes
typedef enum {
    MODE_NORMAL,
    MODE_VISUAL,
    MODE_MOVE
} UIMode;

// Mark storage
#define MAX_MARKS 26  // a-z
static int marks[MAX_MARKS];  // stores clip_index for each mark
static int mark_selected = -1;  // -1 = no mark selected

// Visual mode state
static int visual_start_row = 0;
static int visual_start_col = 0;
static int visual_end_row = 0;
static int visual_end_col = 0;

// Yank buffer
typedef struct {
    int *clip_indices;
    int count;
} YankBuffer;
static YankBuffer yank_buffer = {NULL, 0};

// Current mode
static UIMode current_mode = MODE_NORMAL;

// Convert clip state to color pair
static int state_to_color(ClipState state) {
    switch (state) {
        case CLIP_EMPTY: return COLOR_EMPTY;
        case CLIP_RECORDING: return COLOR_RECORDING;
        case CLIP_LOOPING: return COLOR_LOOPING;
        case CLIP_STOPPED: return COLOR_STOPPED;
        default: return COLOR_EMPTY;
    }
}

// Get clip index from grid position
static int grid_to_clip_index(int row, int col) {
    return row * GRID_COLS + col;
}

// Check if a cell is in the visual selection
static bool is_in_visual_selection(int row, int col) {
    if (current_mode != MODE_VISUAL) return false;
    
    int min_row = (visual_start_row < visual_end_row) ? visual_start_row : visual_end_row;
    int max_row = (visual_start_row > visual_end_row) ? visual_start_row : visual_end_row;
    int min_col = (visual_start_col < visual_end_col) ? visual_start_col : visual_end_col;
    int max_col = (visual_start_col > visual_end_col) ? visual_start_col : visual_end_col;
    
    return (row >= min_row && row <= max_row && col >= min_col && col <= max_col);
}

// Get all clip indices in the visual selection
static int* get_selected_clips(int *count) {
    int min_row = (visual_start_row < visual_end_row) ? visual_start_row : visual_end_row;
    int max_row = (visual_start_row > visual_end_row) ? visual_start_row : visual_end_row;
    int min_col = (visual_start_col < visual_end_col) ? visual_start_col : visual_end_col;
    int max_col = (visual_start_col > visual_end_col) ? visual_start_col : visual_end_col;
    
    int num_rows = max_row - min_row + 1;
    int num_cols = max_col - min_col + 1;
    *count = num_rows * num_cols;
    
    int *clips = (int *)malloc(*count * sizeof(int));
    if (!clips) {
        *count = 0;
        return NULL;
    }
    
    int idx = 0;
    for (int r = min_row; r <= max_row; r++) {
        for (int c = min_col; c <= max_col; c++) {
            clips[idx++] = grid_to_clip_index(r, c);
        }
    }
    
    return clips;
}

// Get all clip indices in a row
static int* get_row_clips(int row, int *count) {
    *count = GRID_COLS;
    int *clips = (int *)malloc(*count * sizeof(int));
    if (!clips) {
        *count = 0;
        return NULL;
    }
    
    for (int c = 0; c < GRID_COLS; c++) {
        clips[c] = grid_to_clip_index(row, c);
    }
    
    return clips;
}

// Delete (reset) clips
static void delete_clips(int *clip_indices, int count) {
    for (int i = 0; i < count; i++) {
        engine_reset_clip(g_engine, clip_indices[i]);
    }
    engine_process_commands(g_engine);
}

// Yank clips (store their indices)
static void yank_clips(int *clip_indices, int count) {
    // Free previous yank buffer
    if (yank_buffer.clip_indices) {
        free(yank_buffer.clip_indices);
    }
    
    yank_buffer.clip_indices = (int *)malloc(count * sizeof(int));
    if (yank_buffer.clip_indices) {
        memcpy(yank_buffer.clip_indices, clip_indices, count * sizeof(int));
        yank_buffer.count = count;
    }
}

// Paste clips (trigger them)
static void paste_clips(void) {
    if (!yank_buffer.clip_indices || yank_buffer.count == 0) return;
    
    // Calculate offset from current position to first yanked clip
    int first_yanked_row = yank_buffer.clip_indices[0] / GRID_COLS;
    int first_yanked_col = yank_buffer.clip_indices[0] % GRID_COLS;
    int row_offset = selected_row - first_yanked_row;
    int col_offset = selected_col - first_yanked_col;
    
    for (int i = 0; i < yank_buffer.count; i++) {
        int orig_row = yank_buffer.clip_indices[i] / GRID_COLS;
        int orig_col = yank_buffer.clip_indices[i] % GRID_COLS;
        
        int new_row = orig_row + row_offset;
        int new_col = orig_col + col_offset;
        
        // Bounds check
        if (new_row >= 0 && new_row < GRID_ROWS && new_col >= 0 && new_col < GRID_COLS) {
            int new_clip_idx = grid_to_clip_index(new_row, new_col);
            // Trigger three times to go from empty -> recording -> looping -> stopped
            engine_trigger_clip(g_engine, new_clip_idx);
            engine_trigger_clip(g_engine, new_clip_idx);
            engine_trigger_clip(g_engine, new_clip_idx);
        }
    }
    engine_process_commands(g_engine);
}

// Set a mark at current position
static void set_mark(char mark_char) {
    if (mark_char >= 'a' && mark_char <= 'z') {
        int idx = mark_char - 'a';
        marks[idx] = grid_to_clip_index(selected_row, selected_col);
    }
}

// Go to a mark
static void go_to_mark(char mark_char) {
    if (mark_char >= 'a' && mark_char <= 'z') {
        int idx = mark_char - 'a';
        int clip_idx = marks[idx];
        if (clip_idx >= 0 && clip_idx < GRID_ROWS * GRID_COLS) {
            selected_row = clip_idx / GRID_COLS;
            selected_col = clip_idx % GRID_COLS;
        }
    }
}

// Play next scene (next row)
static void play_next_scene(void) {
    int next_row = (selected_row + 1) % GRID_ROWS;
    engine_trigger_scene(g_engine, next_row);
    engine_process_commands(g_engine);
    selected_row = next_row;
}

// Play previous scene (previous row)
static void play_prev_scene(void) {
    int prev_row = (selected_row - 1 + GRID_ROWS) % GRID_ROWS;
    engine_trigger_scene(g_engine, prev_row);
    engine_process_commands(g_engine);
    selected_row = prev_row;
}

// Draw a single cell
static void draw_cell(int row, int col, bool selected) {
    int clip_idx = grid_to_clip_index(row, col);
    Clip *clip = &g_engine->clips[clip_idx];
    
    int y = row * CELL_HEIGHT + 1;
    int x = col * CELL_WIDTH + 1;
    
    int color = state_to_color(clip->state);
    if (selected) {
        color = COLOR_SELECTED;
    } else if (current_mode == MODE_VISUAL && is_in_visual_selection(row, col)) {
        // Use a different highlight for visual selection (invert colors)
        color = COLOR_SELECTED;  // Same as selected for now
    }
    
    attron(COLOR_PAIR(color));
    
    // Draw cell border
    for (int dy = 0; dy < CELL_HEIGHT; dy++) {
        for (int dx = 0; dx < CELL_WIDTH; dx++) {
            mvaddch(y + dy, x + dx, ' ');
        }
    }
    
    // Draw clip number
    mvprintw(y + 1, x + 1, "%2d", clip_idx);
    
    // Draw state indicator
    char state_char;
    switch (clip->state) {
        case CLIP_EMPTY: state_char = ' '; break;
        case CLIP_RECORDING: state_char = 'R'; break;
        case CLIP_LOOPING: state_char = 'L'; break;
        case CLIP_STOPPED: state_char = 'S'; break;
        default: state_char = '?'; break;
    }
    mvaddch(y + 1, x + 4, state_char);
    
    attroff(COLOR_PAIR(color));
}

// Draw the entire grid
static void draw_grid(void) {
    clear();
    
    // Draw title
    attron(A_BOLD);
    mvprintw(0, 0, "JACK Looper - 8x8 Clip Grid");
    attroff(A_BOLD);
    
    // Draw cells
    for (int row = 0; row < GRID_ROWS; row++) {
        for (int col = 0; col < GRID_COLS; col++) {
            bool selected = (row == selected_row && col == selected_col);
            draw_cell(row, col, selected);
        }
    }
    
    // Draw status bar
    int clip_idx = grid_to_clip_index(selected_row, selected_col);
    Clip *clip = &g_engine->clips[clip_idx];
    
    mvprintw(GRID_ROWS * CELL_HEIGHT + 1, 0, 
             "Selected: Clip %d | State: %s | Buffer: %zu samples",
             clip_idx, clip_state_to_string(clip->state), clip->buffer_size);
    
    mvprintw(GRID_ROWS * CELL_HEIGHT + 2, 0,
             "Quantize: %s | Threshold: %u",
             quantize_mode_to_string((QuantizeMode)atomic_load(&g_engine->quantize_mode_atomic)),
             (unsigned int)atomic_load(&g_engine->quantize_threshold_atomic));
    
    // Draw mode indicator
    const char *mode_str;
    switch (current_mode) {
        case MODE_NORMAL: mode_str = "NORMAL"; break;
        case MODE_VISUAL: mode_str = "VISUAL"; break;
        case MODE_MOVE:   mode_str = "MOVE";   break;
        default:          mode_str = "?";       break;
    }
    mvprintw(GRID_ROWS * CELL_HEIGHT + 3, 0,
             "Mode: %s", mode_str);
    
    // Draw help if active
    if (show_help) {
        attron(COLOR_PAIR(COLOR_HELP));
        mvprintw(GRID_ROWS * CELL_HEIGHT + 4, 0, 
                 "=== Help ===");
        mvprintw(GRID_ROWS * CELL_HEIGHT + 5, 0,
                 "h/j/k/l  - Navigate grid (left/down/up/right)");
        mvprintw(GRID_ROWS * CELL_HEIGHT + 6, 0,
                 "t        - Trigger selected clip");
        mvprintw(GRID_ROWS * CELL_HEIGHT + 7, 0,
                 "r        - Reset selected clip");
        mvprintw(GRID_ROWS * CELL_HEIGHT + 8, 0,
                 "s        - Trigger scene (current row)");
        mvprintw(GRID_ROWS * CELL_HEIGHT + 9, 0,
                 "q        - Toggle quantize mode (off/beat/bar)");
        mvprintw(GRID_ROWS * CELL_HEIGHT + 10, 0,
                 "T        - Set quantize threshold");
        mvprintw(GRID_ROWS * CELL_HEIGHT + 11, 0,
                 "x        - Reset transport");
        mvprintw(GRID_ROWS * CELL_HEIGHT + 12, 0,
                 "?        - Toggle help");
        mvprintw(GRID_ROWS * CELL_HEIGHT + 13, 0,
                 "Esc/q   - Quit");
        attroff(COLOR_PAIR(COLOR_HELP));
    }
    
    refresh();
}

// Handle command mode input (after pressing ':')
// Returns true if the application should quit
static bool handle_command_mode(void) {
    char cmd_buffer[256];
    int cmd_pos = 0;
    memset(cmd_buffer, 0, sizeof(cmd_buffer));
    
    // Save current nodelay state and force blocking input
    int prev_nodelay = nodelay(stdscr, FALSE);
    
    // Show command prompt
    mvprintw(LINES - 1, 0, ":");
    clrtoeol();
    refresh();
    
    while (1) {
        int ch = getch();
        // Do NOT break on ERR; getch() will block now
        
        if (ch == '\n' || ch == '\r') {
            // Execute command
            cmd_buffer[cmd_pos] = '\0';
            
            // Clear command line
            mvprintw(LINES - 1, 0, "                ");
            refresh();
            
            // Parse and execute command
            if (strcmp(cmd_buffer, "q") == 0) {
                // Restore previous nodelay state before returning
                nodelay(stdscr, prev_nodelay);
                return true;  // Quit
            }
            // Add more commands here as needed
            
            // Restore previous nodelay state before returning
            nodelay(stdscr, prev_nodelay);
            return false;  // Don't quit
        } else if (ch == 27) {  // Escape - cancel command mode
            mvprintw(LINES - 1, 0, "                ");
            refresh();
            nodelay(stdscr, prev_nodelay);
            return false;
        } else if (ch == KEY_BACKSPACE || ch == 127) {  // Backspace
            if (cmd_pos > 0) {
                cmd_pos--;
                cmd_buffer[cmd_pos] = '\0';
                mvprintw(LINES - 1, 0, ":%s ", cmd_buffer);
                refresh();
            }
        } else if (cmd_pos < (int)sizeof(cmd_buffer) - 1) {
            cmd_buffer[cmd_pos++] = (char)ch;
            cmd_buffer[cmd_pos] = '\0';
            mvprintw(LINES - 1, 0, ":%s", cmd_buffer);
            refresh();
        }
    }
    
    // Should never reach here, but restore just in case
    nodelay(stdscr, prev_nodelay);
    return false;
}

static void handle_sigint(int sig) {
    (void)sig;
    tui_cleanup();
    _Exit(1);
}

void tui_init(Engine *engine) {
    g_engine = engine;
    
    // Initialize ncurses
    initscr();
    cbreak();
    noecho();
    keypad(stdscr, TRUE);
    curs_set(0);  // Hide cursor
    
    // Initialize colors
    if (has_colors()) {
        start_color();
        
        // Define color pairs
        init_pair(COLOR_EMPTY, COLOR_WHITE, COLOR_BLACK);
        init_pair(COLOR_RECORDING, COLOR_RED, COLOR_BLACK);
        init_pair(COLOR_LOOPING, COLOR_GREEN, COLOR_BLACK);
        init_pair(COLOR_STOPPED, COLOR_BLUE, COLOR_BLACK);
        init_pair(COLOR_SELECTED, COLOR_BLACK, COLOR_CYAN);
        init_pair(COLOR_HELP, COLOR_CYAN, COLOR_BLACK);
    }
}

void tui_run(Engine *engine) {
    if (!engine) return;
    
    g_engine = engine;
    
    // Initialize marks
    for (int i = 0; i < MAX_MARKS; i++) {
        marks[i] = -1;
    }
    
    draw_grid();
    
    while (1) {
        int ch = getch();
        if (ch == ERR) {
            break;
        }
        
        // Handle mode-specific input
        if (current_mode == MODE_MOVE) {
            switch (ch) {
                case 'h':
                case KEY_LEFT:
                    selected_col = (selected_col - 1 + GRID_COLS) % GRID_COLS;
                    break;
                    
                case 'j':
                case KEY_DOWN:
                    selected_row = (selected_row + 1) % GRID_ROWS;
                    break;
                    
                case 'k':
                case KEY_UP:
                    selected_row = (selected_row - 1 + GRID_ROWS) % GRID_ROWS;
                    break;
                    
                case 'l':
                case KEY_RIGHT:
                    selected_col = (selected_col + 1) % GRID_COLS;
                    break;
                    
                case '\n':
                case '\r':
                case 27:  // Escape
                    current_mode = MODE_NORMAL;
                    break;
                    
                default:
                    break;
            }
            
            draw_grid();
            continue;
        }
        
        if (current_mode == MODE_VISUAL) {
            switch (ch) {
                case 'h':
                case KEY_LEFT:
                    visual_end_col = (visual_end_col - 1 + GRID_COLS) % GRID_COLS;
                    break;
                    
                case 'j':
                case KEY_DOWN:
                    visual_end_row = (visual_end_row + 1) % GRID_ROWS;
                    break;
                    
                case 'k':
                case KEY_UP:
                    visual_end_row = (visual_end_row - 1 + GRID_ROWS) % GRID_ROWS;
                    break;
                    
                case 'l':
                case KEY_RIGHT:
                    visual_end_col = (visual_end_col + 1) % GRID_COLS;
                    break;
                    
                case 'd': {
                    int count;
                    int *clips = get_selected_clips(&count);
                    if (clips) {
                        delete_clips(clips, count);
                        free(clips);
                    }
                    current_mode = MODE_NORMAL;
                    break;
                }
                    
                case 'y': {
                    int count;
                    int *clips = get_selected_clips(&count);
                    if (clips) {
                        yank_clips(clips, count);
                        free(clips);
                    }
                    current_mode = MODE_NORMAL;
                    break;
                }
                    
                case 27:  // Escape
                    current_mode = MODE_NORMAL;
                    break;
                    
                default:
                    break;
            }
            
            draw_grid();
            continue;
        }
        
        // Normal mode
        switch (ch) {
            case 'h':
            case KEY_LEFT:
                selected_col = (selected_col - 1 + GRID_COLS) % GRID_COLS;
                break;
                
            case 'j':
            case KEY_DOWN:
                selected_row = (selected_row + 1) % GRID_ROWS;
                break;
                
            case 'k':
            case KEY_UP:
                selected_row = (selected_row - 1 + GRID_ROWS) % GRID_ROWS;
                break;
                
            case 'l':
            case KEY_RIGHT:
                selected_col = (selected_col + 1) % GRID_COLS;
                break;
                
            case 't': {
                int clip_idx = grid_to_clip_index(selected_row, selected_col);
                engine_trigger_clip(engine, clip_idx);
                engine_process_commands(engine);
                break;
            }
                
            case 'd': {
                // Delete (reset) current clip
                int clip_idx = grid_to_clip_index(selected_row, selected_col);
                engine_reset_clip(engine, clip_idx);
                engine_process_commands(engine);
                break;
            }
                
            case 'y': {
                // Yank current clip
                int clip_idx = grid_to_clip_index(selected_row, selected_col);
                yank_clips(&clip_idx, 1);
                break;
            }
                
            case 'p': {
                // Paste
                paste_clips();
                break;
            }
                
            case 's': {
                // Trigger scene for current row
                engine_trigger_scene(engine, selected_row);
                engine_process_commands(engine);
                break;
            }
                
            case 'q': {
                // Cycle quantize mode
                QuantizeMode modes[] = {QUANTIZE_OFF, QUANTIZE_BEAT, QUANTIZE_BAR};
                int num_modes = sizeof(modes) / sizeof(modes[0]);
                int current = 0;
                for (int i = 0; i < num_modes; i++) {
                    if (engine->quantize_mode == modes[i]) {
                        current = i;
                        break;
                    }
                }
                QuantizeMode next = modes[(current + 1) % num_modes];
                engine_set_quantize_mode(engine, next);
                engine_process_commands(engine);
                break;
            }
                
            case 'T': {
                // Toggle threshold between 0 and 1000
                if (engine->quantize_threshold == 0) {
                    engine_set_quantize_threshold(engine, 1000);
                } else {
                    engine_set_quantize_threshold(engine, 0);
                }
                engine_process_commands(engine);
                break;
            }
                
            case 'x':
                engine_reset_transport(engine);
                engine_process_commands(engine);
                break;
                
            case ':': {
                bool should_quit = handle_command_mode();
                if (should_quit) {
                    return;
                }
                break;
            }
                
            case '?':
                show_help = !show_help;
                break;
                
            case 'v': {
                // Enter visual mode
                current_mode = MODE_VISUAL;
                visual_start_row = selected_row;
                visual_start_col = selected_col;
                visual_end_row = selected_row;
                visual_end_col = selected_col;
                break;
            }
                
            case 'V': {
                // Select entire line
                current_mode = MODE_VISUAL;
                visual_start_row = selected_row;
                visual_start_col = 0;
                visual_end_row = selected_row;
                visual_end_col = GRID_COLS - 1;
                break;
            }
                
            case 'm': {
                // Enter move mode
                current_mode = MODE_MOVE;
                break;
            }
                
            case 'N': {
                // Play next scene
                play_next_scene();
                break;
            }
                
            case 'P': {
                // Play previous scene
                play_prev_scene();
                break;
            }
                
            case '\'': {
                // Go to mark - wait for next character
                nodelay(stdscr, FALSE);
                int mark_ch = getch();
                nodelay(stdscr, TRUE);
                if (mark_ch != ERR) {
                    go_to_mark((char)mark_ch);
                }
                break;
            }
                
            case 27:  // Escape key
            case 'Q':
                return;
                
            default:
                // Check for mark setting (m<char>)
                if (ch == 'm') {
                    // Wait for next character
                    nodelay(stdscr, FALSE);
                    int mark_ch = getch();
                    nodelay(stdscr, TRUE);
                    if (mark_ch != ERR) {
                        set_mark((char)mark_ch);
                    }
                }
                break;
        }
        
        draw_grid();
    }
}

void tui_cleanup(void) {
    // Free yank buffer
    if (yank_buffer.clip_indices) {
        free(yank_buffer.clip_indices);
        yank_buffer.clip_indices = NULL;
        yank_buffer.count = 0;
    }
    
    // Restore terminal settings
    curs_set(1);
    endwin();
    
    // Reset signal handler to default
    signal(SIGINT, SIG_DFL);
}