#!/usr/bin/env python3
"""Vim-like terminal editor with Dracula Pro syntax highlighting."""

import curses
import sys
import os
import re
from pathlib import Path

try:
    from pygments import lex
    from pygments.lexers import PythonLexer, get_lexer_for_filename, TextLexer
    from pygments.token import (
        Token, Comment, String, Number, Keyword, Name, Operator, Punctuation,
    )
    HAS_PYGMENTS = True
except ImportError:
    HAS_PYGMENTS = False

# ── Dracula Pro palette mapped to curses (closest 256-color matches) ─────────
# bg: #22212C  fg: #F8F8F2
# comment: #7970A9  string: #19F9D8  number: #BF9EEE
# keyword: #FF80BF  function: #62E884  operator: #FF9580
# class: #9580FF  builtin: #80FFEA  error: #FF5555

DRACULA = {
    "bg":       236,   # ~#303030
    "fg":       255,   # white
    "comment":  97,    # purple-grey
    "string":   43,    # cyan/teal
    "number":   141,   # light purple
    "keyword":  212,   # pink
    "function": 84,    # green
    "operator": 209,   # orange
    "class":    141,   # purple
    "builtin":  87,    # bright cyan
    "error":    196,   # red
    "lineno":   240,   # grey
    "status":   141,   # purple
    "statusfg": 16,    # black
    "visual":   60,    # selection highlight bg
    "cursor":   255,   # cursor line
    "search":   220,   # yellow highlight
}

# Map pygments token types to our color names
TOKEN_MAP = {
    Comment:              "comment",
    Comment.Single:       "comment",
    Comment.Multiline:    "comment",
    Comment.Hashbang:     "comment",
    String:               "string",
    String.Single:        "string",
    String.Double:        "string",
    String.Doc:           "string",
    String.Escape:        "operator",
    String.Interpol:      "operator",
    String.Affix:         "keyword",
    Number:               "number",
    Number.Integer:       "number",
    Number.Float:         "number",
    Number.Hex:           "number",
    Number.Oct:           "number",
    Number.Bin:           "number",
    Keyword:              "keyword",
    Keyword.Constant:     "keyword",
    Keyword.Declaration:  "keyword",
    Keyword.Namespace:    "keyword",
    Keyword.Pseudo:       "keyword",
    Keyword.Reserved:     "keyword",
    Keyword.Type:         "class",
    Name.Function:        "function",
    Name.Function.Magic:  "function",
    Name.Class:           "class",
    Name.Decorator:       "function",
    Name.Builtin:         "builtin",
    Name.Builtin.Pseudo:  "builtin",
    Name.Exception:       "error",
    Operator:             "operator",
    Operator.Word:        "keyword",
    Punctuation:          "fg",
    Name:                 "fg",
    Token.Text:           "fg",
}


def token_color(ttype):
    """Walk up the token type hierarchy to find a color."""
    while ttype:
        if ttype in TOKEN_MAP:
            return TOKEN_MAP[ttype]
        ttype = ttype.parent
    return "fg"


class Editor:
    NORMAL = "NORMAL"
    INSERT = "INSERT"
    VISUAL = "VISUAL"
    VISUAL_LINE = "V-LINE"
    COMMAND = "COMMAND"
    SEARCH = "SEARCH"

    def __init__(self, stdscr, filename=None):
        self.stdscr = stdscr
        self.filename = filename
        self.lines = [""]
        self.cx = 0          # cursor x in line
        self.cy = 0          # cursor y (line index)
        self.scroll_y = 0
        self.scroll_x = 0
        self.mode = self.NORMAL
        self.cmd_buf = ""
        self.search_query = ""
        self.search_dir = 1   # 1=forward, -1=backward
        self.search_matches = []
        self.message = ""
        self.pending = ""     # pending normal-mode key sequence
        self.visual_start = (0, 0)  # (y, x)
        self.modified = False
        self.undo_stack = []
        self.redo_stack = []
        self.last_save_state = None
        self.yank_buf = []
        self.yank_linewise = False
        self.repeat_count = ""
        self.last_search = ""
        self.gutter_w = 4
        # folding: set of line indices that are fold headers (collapsed)
        self.folds = {}  # {start_line: end_line} - collapsed ranges
        self.fold_all = False

        # curses setup
        curses.curs_set(1)
        curses.start_color()
        curses.use_default_colors()
        self._init_colors()
        self.stdscr.timeout(-1)

        if filename and Path(filename).exists():
            self.load_file(filename)

        self.save_undo()

    def _init_colors(self):
        """Set up 256-color pairs for Dracula Pro."""
        self.color_pairs = {}
        pair_id = 1
        bg = DRACULA["bg"]
        for name, fg in DRACULA.items():
            if name == "bg":
                continue
            curses.init_pair(pair_id, fg, bg)
            self.color_pairs[name] = pair_id
            pair_id += 1
        # visual selection
        curses.init_pair(pair_id, DRACULA["fg"], DRACULA["visual"])
        self.color_pairs["visual_sel"] = pair_id
        pair_id += 1
        # search highlight
        curses.init_pair(pair_id, 16, DRACULA["search"])
        self.color_pairs["search_hl"] = pair_id
        pair_id += 1
        # status bar
        curses.init_pair(pair_id, DRACULA["statusfg"], DRACULA["status"])
        self.color_pairs["statusbar"] = pair_id
        pair_id += 1
        # bg fill
        curses.init_pair(pair_id, DRACULA["fg"], DRACULA["bg"])
        self.color_pairs["default"] = pair_id

    def cp(self, name):
        return curses.color_pair(self.color_pairs.get(name, 0))

    # ── File I/O ─────────────────────────────────────────────────────────
    def load_file(self, path):
        try:
            with open(path, "r") as f:
                text = f.read()
            self.lines = text.split("\n")
            if self.lines and self.lines[-1] == "":
                self.lines = self.lines[:-1] or [""]
        except Exception as e:
            self.message = f"Error: {e}"
            self.lines = [""]

    def save_file(self):
        if not self.filename:
            self.message = "No filename"
            return False
        try:
            with open(self.filename, "w") as f:
                f.write("\n".join(self.lines) + "\n")
            self.modified = False
            self.message = f'"{self.filename}" written, {len(self.lines)}L'
            return True
        except Exception as e:
            self.message = f"Error: {e}"
            return False

    # ── Undo/Redo ────────────────────────────────────────────────────────
    def save_undo(self):
        state = ([l for l in self.lines], self.cx, self.cy)
        if self.undo_stack and self.undo_stack[-1][0] == state[0]:
            return
        self.undo_stack.append(state)
        if len(self.undo_stack) > 500:
            self.undo_stack.pop(0)
        self.redo_stack.clear()

    def undo(self):
        if len(self.undo_stack) <= 1:
            self.message = "Already at oldest change"
            return
        self.redo_stack.append(self.undo_stack.pop())
        lines, cx, cy = self.undo_stack[-1]
        self.lines = [l for l in lines]
        self.cx = cx
        self.cy = cy
        self.modified = True

    def redo(self):
        if not self.redo_stack:
            self.message = "Already at newest change"
            return
        state = self.redo_stack.pop()
        self.undo_stack.append(state)
        self.lines = [l for l in state[0]]
        self.cx = state[1]
        self.cy = state[2]
        self.modified = True

    # ── Syntax highlighting ──────────────────────────────────────────────
    def get_lexer(self):
        if not HAS_PYGMENTS:
            return None
        if self.filename:
            try:
                return get_lexer_for_filename(self.filename)
            except Exception:
                pass
        return PythonLexer()

    def highlight_line(self, text, line_idx):
        """Return list of (char, color_name) for a line. Uses full-file lex for accuracy."""
        if not HAS_PYGMENTS:
            return [(ch, "fg") for ch in text]

        if not hasattr(self, "_hl_cache_id") or self._hl_cache_id != id(self.lines):
            self._rebuild_highlight_cache()

        if line_idx < len(self._hl_lines):
            return self._hl_lines[line_idx]
        return [(ch, "fg") for ch in text]

    def _rebuild_highlight_cache(self):
        lexer = self.get_lexer()
        if not lexer:
            self._hl_cache_id = id(self.lines)
            self._hl_lines = []
            return

        full_text = "\n".join(self.lines)
        self._hl_lines = [[] for _ in self.lines]
        line_idx = 0
        col = 0

        for ttype, value in lex(full_text, lexer):
            color = token_color(ttype)
            for ch in value:
                if ch == "\n":
                    line_idx += 1
                    col = 0
                else:
                    if line_idx < len(self._hl_lines):
                        self._hl_lines[line_idx].append((ch, color))
                    col += 1

        self._hl_cache_id = id(self.lines)

    # ── Drawing ──────────────────────────────────────────────────────────
    def draw(self):
        self.stdscr.erase()
        h, w = self.stdscr.getmaxyx()
        text_h = h - 2   # status + cmd line
        self.gutter_w = max(4, len(str(len(self.lines))) + 2)
        text_w = w - self.gutter_w

        # fill bg
        bg = self.cp("default")
        for row in range(h):
            try:
                self.stdscr.addnstr(row, 0, " " * w, w, bg)
            except curses.error:
                pass

        # ensure cursor not inside fold
        self.ensure_cursor_visible()

        # build visible line list
        vis = self.visible_lines()

        # find cursor position in visible lines
        cur_vis_idx = 0
        for vi, li in enumerate(vis):
            if li == self.cy:
                cur_vis_idx = vi
                break

        # scroll only to keep cursor on screen (no auto-unscroll)
        if cur_vis_idx < self.scroll_y:
            self.scroll_y = cur_vis_idx
        if cur_vis_idx >= self.scroll_y + text_h:
            self.scroll_y = cur_vis_idx - text_h + 1
        if self.cx < self.scroll_x:
            self.scroll_x = self.cx
        if self.cx >= self.scroll_x + text_w:
            self.scroll_x = self.cx - text_w + 1

        # visual selection range
        in_visual = self.mode in (self.VISUAL, self.VISUAL_LINE)
        vis_start = vis_end = None
        if in_visual:
            sy, sx = self.visual_start
            ey, ex = self.cy, self.cx
            if (sy, sx) > (ey, ex):
                sy, sx, ey, ex = ey, ex, sy, sx
            vis_start = (sy, sx)
            vis_end = (ey, ex)

        # draw lines
        for row in range(text_h):
            vis_idx = self.scroll_y + row
            if vis_idx >= len(vis):
                # tilde for empty lines
                try:
                    self.stdscr.addnstr(row, 0, " ~", self.gutter_w, self.cp("comment"))
                except curses.error:
                    pass
                continue

            line_idx = vis[vis_idx]

            # fold marker
            fold_marker = ""
            if line_idx in self.folds:
                fold_end = self.folds[line_idx]
                fold_marker = f" +-- {fold_end - line_idx} lines "

            # line number
            lnum = str(line_idx + 1).rjust(self.gutter_w - 1) + " "
            try:
                self.stdscr.addnstr(row, 0, lnum, self.gutter_w, self.cp("lineno"))
            except curses.error:
                pass

            line = self.lines[line_idx]
            hl = self.highlight_line(line, line_idx)

            for col in range(text_w):
                char_idx = self.scroll_x + col
                if char_idx >= len(line):
                    break

                ch = line[char_idx]
                if char_idx < len(hl):
                    color_name = hl[char_idx][1]
                else:
                    color_name = "fg"

                attr = self.cp(color_name)

                # visual highlight
                if in_visual:
                    if self.mode == self.VISUAL_LINE:
                        if vis_start[0] <= line_idx <= vis_end[0]:
                            attr = self.cp("visual_sel")
                    else:
                        pos = (line_idx, char_idx)
                        if vis_start <= pos <= vis_end:
                            attr = self.cp("visual_sel")

                # search highlight
                if self.last_search:
                    slen = len(self.last_search)
                    if slen > 0 and char_idx <= len(line) - slen:
                        segment = line[char_idx:char_idx + slen].lower()
                        if segment == self.last_search.lower() and char_idx == line.lower().find(self.last_search.lower(), char_idx):
                            attr = self.cp("search_hl")

                try:
                    self.stdscr.addstr(row, self.gutter_w + col, ch, attr)
                except curses.error:
                    pass

            # draw fold marker after line content
            if fold_marker:
                col_start = min(len(line) - self.scroll_x, text_w)
                if col_start < 0:
                    col_start = 0
                try:
                    self.stdscr.addnstr(row, self.gutter_w + col_start, fold_marker,
                                        text_w - col_start, self.cp("comment") | curses.A_BOLD)
                except curses.error:
                    pass

        # status bar
        status_row = h - 2
        mode_str = f" {self.mode} "
        fname = self.filename or "[No Name]"
        mod = " [+]" if self.modified else ""
        pos = f" {self.cy + 1}:{self.cx + 1} "
        pct = f" {int((self.cy + 1) / max(len(self.lines), 1) * 100)}% "

        left = f"{mode_str} {fname}{mod}"
        right = f"{pos}{pct}"
        gap = w - len(left) - len(right)
        if gap < 0:
            gap = 0
        bar = left + " " * gap + right

        try:
            self.stdscr.addnstr(status_row, 0, bar, w, self.cp("statusbar") | curses.A_BOLD)
        except curses.error:
            pass

        # command / message line
        cmd_row = h - 1
        if self.mode == self.COMMAND:
            cmd_text = ":" + self.cmd_buf
        elif self.mode == self.SEARCH:
            ch = "/" if self.search_dir == 1 else "?"
            cmd_text = ch + self.search_query
        elif self.message:
            cmd_text = self.message
        else:
            cmd_text = ""
        try:
            self.stdscr.addnstr(cmd_row, 0, cmd_text[:w-1], w-1, self.cp("fg"))
        except curses.error:
            pass

        # position cursor
        scr_y = cur_vis_idx - self.scroll_y
        scr_x = self.gutter_w + self.cx - self.scroll_x
        if 0 <= scr_y < text_h and 0 <= scr_x < w:
            if self.mode != self.INSERT:
                # block cursor: highlight char under cursor with reverse video
                if self.cy < len(self.lines) and self.cx < len(self.lines[self.cy]):
                    ch = self.lines[self.cy][self.cx]
                else:
                    ch = " "
                try:
                    self.stdscr.addstr(scr_y, scr_x, ch, curses.A_REVERSE | self.cp("fg"))
                except curses.error:
                    pass
                curses.curs_set(0)
            else:
                try:
                    self.stdscr.move(scr_y, scr_x)
                except curses.error:
                    pass
                curses.curs_set(2)

        self.stdscr.refresh()

    # ── Cursor helpers ───────────────────────────────────────────────────
    def clamp_cursor(self):
        self.cy = max(0, min(self.cy, len(self.lines) - 1))
        line_len = len(self.lines[self.cy])
        if self.mode == self.INSERT:
            self.cx = max(0, min(self.cx, line_len))
        else:
            self.cx = max(0, min(self.cx, line_len - 1)) if line_len > 0 else 0

    def cur_line(self):
        return self.lines[self.cy]

    # ── Folding ──────────────────────────────────────────────────────────
    def _get_indent_level(self, line_idx):
        line = self.lines[line_idx]
        stripped = line.lstrip()
        if not stripped:
            return -1  # blank line
        return len(line) - len(stripped)

    def _find_fold_end(self, start):
        """Find the end of a fold block starting at `start` (indent-based)."""
        base_indent = self._get_indent_level(start)
        if base_indent < 0:
            return start
        end = start
        for i in range(start + 1, len(self.lines)):
            lvl = self._get_indent_level(i)
            if lvl == -1:  # blank line, continue
                end = i
                continue
            if lvl > base_indent:
                end = i
            else:
                break
        return end if end > start else start

    def fold_close(self, line_idx):
        """Collapse the fold at line_idx."""
        end = self._find_fold_end(line_idx)
        if end > line_idx:
            self.folds[line_idx] = end

    def fold_open(self, line_idx):
        """Open the fold at line_idx."""
        self.folds.pop(line_idx, None)

    def fold_all_close(self):
        """zM: close all top-level folds."""
        self.folds.clear()
        i = 0
        while i < len(self.lines):
            line = self.lines[i]
            stripped = line.rstrip()
            indent = self._get_indent_level(i)
            if indent == 0 and stripped and stripped.endswith(":"):
                end = self._find_fold_end(i)
                if end > i:
                    self.folds[i] = end
                    i = end + 1
                    continue
            i += 1

    def fold_all_open(self):
        """zO: open all folds."""
        self.folds.clear()

    def is_line_hidden(self, line_idx):
        """Check if a line is inside a collapsed fold."""
        for start, end in self.folds.items():
            if start < line_idx <= end:
                return True
        return False

    def visible_lines(self):
        """Return list of visible line indices (not hidden by folds)."""
        result = []
        for i in range(len(self.lines)):
            if not self.is_line_hidden(i):
                result.append(i)
        return result

    def ensure_cursor_visible(self):
        """If cursor is inside a fold, move it to the fold header."""
        for start, end in self.folds.items():
            if start < self.cy <= end:
                self.cy = start
                break

    # ── Paragraph motions ────────────────────────────────────────────────
    def para_forward(self):
        """Jump to next blank line (or EOF)."""
        y = self.cy + 1
        # skip current non-blank
        while y < len(self.lines) and self.lines[y].strip():
            y += 1
        # skip blank
        while y < len(self.lines) and not self.lines[y].strip():
            y += 1
        self.cy = min(y, len(self.lines) - 1)
        self.cx = 0

    def para_backward(self):
        """Jump to previous blank line (or BOF)."""
        y = self.cy - 1
        # skip current non-blank
        while y > 0 and self.lines[y].strip():
            y -= 1
        # skip blank
        while y > 0 and not self.lines[y].strip():
            y -= 1
        self.cy = max(y, 0)
        self.cx = 0

    # ── Word motions ─────────────────────────────────────────────────────
    def word_forward(self):
        line = self.cur_line()
        x = self.cx
        if x >= len(line) - 1:
            if self.cy < len(self.lines) - 1:
                self.cy += 1
                self.cx = 0
                # skip leading whitespace
                line = self.cur_line()
                while self.cx < len(line) and line[self.cx] == " ":
                    self.cx += 1
            return
        # skip current word
        while x < len(line) and line[x] not in " \t":
            x += 1
        # skip whitespace
        while x < len(line) and line[x] in " \t":
            x += 1
        self.cx = x

    def word_backward(self):
        line = self.cur_line()
        x = self.cx
        if x <= 0:
            if self.cy > 0:
                self.cy -= 1
                self.cx = max(0, len(self.lines[self.cy]) - 1)
            return
        x -= 1
        line = self.cur_line()
        while x > 0 and line[x] in " \t":
            x -= 1
        while x > 0 and line[x - 1] not in " \t":
            x -= 1
        self.cx = x

    def word_end(self):
        line = self.cur_line()
        x = self.cx + 1
        if x >= len(line):
            if self.cy < len(self.lines) - 1:
                self.cy += 1
                line = self.lines[self.cy]
                x = 0
                while x < len(line) and line[x] in " \t":
                    x += 1
        while x < len(line) - 1 and line[x + 1] not in " \t":
            x += 1
        self.cx = min(x, max(0, len(self.cur_line()) - 1))

    # ── Search ───────────────────────────────────────────────────────────
    def do_search(self, direction=None):
        if not self.last_search:
            return
        d = direction if direction else self.search_dir
        total = len(self.lines)
        y, x = self.cy, self.cx + d
        for _ in range(total + 1):
            if y < 0:
                y = total - 1
                x = max(0, len(self.lines[y]) - 1)
            elif y >= total:
                y = 0
                x = 0
            line = self.lines[y]
            idx = line.lower().find(self.last_search.lower(), max(0, x))
            if idx >= 0 and not (y == self.cy and idx == self.cx):
                self.cy = y
                self.cx = idx
                self.clamp_cursor()
                return
            y += d
            x = 0 if d == 1 else max(0, len(self.lines[max(0, min(y, total-1))]) - 1)
        self.message = "Pattern not found"

    # ── Delete operations ────────────────────────────────────────────────
    def delete_char(self):
        line = self.cur_line()
        if self.cx < len(line):
            self.lines[self.cy] = line[:self.cx] + line[self.cx + 1:]
            self.modified = True
        elif self.cy < len(self.lines) - 1:
            self.lines[self.cy] = line + self.lines[self.cy + 1]
            self.lines.pop(self.cy + 1)
            self.modified = True

    def delete_line(self, count=1):
        count = min(count, len(self.lines) - self.cy)
        deleted = self.lines[self.cy:self.cy + count]
        self.yank_buf = deleted
        self.yank_linewise = True
        del self.lines[self.cy:self.cy + count]
        if not self.lines:
            self.lines = [""]
        self.clamp_cursor()
        self.modified = True
        return deleted

    def delete_to_eol(self):
        line = self.cur_line()
        deleted = line[self.cx:]
        self.lines[self.cy] = line[:self.cx]
        self.yank_buf = [deleted]
        self.yank_linewise = False
        self.clamp_cursor()
        self.modified = True

    # ── Yank / Paste ─────────────────────────────────────────────────────
    def yank_line(self, count=1):
        count = min(count, len(self.lines) - self.cy)
        self.yank_buf = self.lines[self.cy:self.cy + count]
        self.yank_linewise = True
        self.message = f"{count} line(s) yanked"

    def paste_after(self):
        if not self.yank_buf:
            return
        self.save_undo()
        if self.yank_linewise:
            idx = self.cy + 1
            for l in self.yank_buf:
                self.lines.insert(idx, l)
                idx += 1
            self.cy += 1
            self.cx = 0
        else:
            line = self.cur_line()
            text = self.yank_buf[0] if self.yank_buf else ""
            self.lines[self.cy] = line[:self.cx + 1] + text + line[self.cx + 1:]
            self.cx += len(text)
        self.modified = True

    def paste_before(self):
        if not self.yank_buf:
            return
        self.save_undo()
        if self.yank_linewise:
            idx = self.cy
            for l in self.yank_buf:
                self.lines.insert(idx, l)
                idx += 1
            self.cx = 0
        else:
            line = self.cur_line()
            text = self.yank_buf[0] if self.yank_buf else ""
            self.lines[self.cy] = line[:self.cx] + text + line[self.cx:]
        self.modified = True

    # ── Visual mode ops ──────────────────────────────────────────────────
    def visual_delete(self):
        sy, sx = self.visual_start
        ey, ex = self.cy, self.cx
        if (sy, sx) > (ey, ex):
            sy, sx, ey, ex = ey, ex, sy, sx

        if self.mode == self.VISUAL_LINE:
            self.yank_buf = self.lines[sy:ey + 1]
            self.yank_linewise = True
            del self.lines[sy:ey + 1]
            if not self.lines:
                self.lines = [""]
            self.cy = min(sy, len(self.lines) - 1)
            self.cx = 0
        else:
            if sy == ey:
                line = self.lines[sy]
                deleted = line[sx:ex + 1]
                self.lines[sy] = line[:sx] + line[ex + 1:]
                self.yank_buf = [deleted]
            else:
                first = self.lines[sy][:sx]
                last = self.lines[ey][ex + 1:]
                mid_lines = self.lines[sy:ey + 1]
                self.yank_buf = mid_lines
                self.lines[sy] = first + last
                del self.lines[sy + 1:ey + 1]
            self.yank_linewise = False
            self.cy = sy
            self.cx = sx
        self.clamp_cursor()
        self.modified = True
        self.mode = self.NORMAL

    def visual_yank(self):
        sy, sx = self.visual_start
        ey, ex = self.cy, self.cx
        if (sy, sx) > (ey, ex):
            sy, sx, ey, ex = ey, ex, sy, sx

        if self.mode == self.VISUAL_LINE:
            self.yank_buf = self.lines[sy:ey + 1]
            self.yank_linewise = True
        else:
            if sy == ey:
                self.yank_buf = [self.lines[sy][sx:ex + 1]]
            else:
                self.yank_buf = self.lines[sy:ey + 1]
            self.yank_linewise = False
        self.cy = sy
        self.cx = sx
        self.mode = self.NORMAL
        self.message = "Yanked"

    # ── Indent ───────────────────────────────────────────────────────────
    def indent_line(self, line_idx, direction=1):
        if direction == 1:
            self.lines[line_idx] = "    " + self.lines[line_idx]
        else:
            line = self.lines[line_idx]
            if line.startswith("    "):
                self.lines[line_idx] = line[4:]
            elif line.startswith("\t"):
                self.lines[line_idx] = line[1:]
            else:
                self.lines[line_idx] = line.lstrip(" ")
        self.modified = True

    # ── Command mode ─────────────────────────────────────────────────────
    def exec_command(self):
        cmd = self.cmd_buf.strip()
        self.cmd_buf = ""
        self.mode = self.NORMAL

        if cmd in ("q", "quit"):
            if self.modified:
                self.message = "No write since last change (use :q! to override)"
            else:
                return "quit"
        elif cmd in ("q!", "quit!"):
            return "quit"
        elif cmd in ("w", "write"):
            self.save_file()
        elif cmd in ("wq", "x"):
            if self.save_file():
                return "quit"
        elif cmd.startswith("w "):
            self.filename = cmd[2:].strip()
            self.save_file()
        elif cmd.startswith("e "):
            fname = cmd[2:].strip()
            self.filename = fname
            self.load_file(fname)
            self.save_undo()
        elif cmd == "noh" or cmd == "nohlsearch":
            self.last_search = ""
        elif re.match(r"^\d+$", cmd):
            line = int(cmd) - 1
            self.cy = max(0, min(line, len(self.lines) - 1))
            self.cx = 0
            self.clamp_cursor()
        elif cmd.startswith("s/") or cmd.startswith("%s/"):
            self._substitute(cmd)
        else:
            self.message = f"Unknown command: {cmd}"
        return None

    def _substitute(self, cmd):
        """Handle :s/pat/rep/ and :%s/pat/rep/g"""
        global_scope = cmd.startswith("%")
        if global_scope:
            cmd = cmd[1:]
        parts = cmd.split("/")
        if len(parts) < 3:
            self.message = "Invalid substitute"
            return
        pattern = parts[1]
        replacement = parts[2]
        flags = parts[3] if len(parts) > 3 else ""
        count_flag = 0 if "g" in flags else 1

        self.save_undo()
        total = 0
        lines_range = range(len(self.lines)) if global_scope else range(self.cy, self.cy + 1)
        for i in lines_range:
            new_line, n = re.subn(pattern, replacement, self.lines[i], count=count_flag)
            if n > 0:
                self.lines[i] = new_line
                total += n
        if total:
            self.modified = True
            self.message = f"{total} substitution(s)"
        else:
            self.message = "Pattern not found"

    # ── Get count prefix ─────────────────────────────────────────────────
    def get_count(self):
        if self.repeat_count:
            n = int(self.repeat_count)
            self.repeat_count = ""
            return n
        return 1

    # ── Normal mode ──────────────────────────────────────────────────────
    def handle_normal(self, key):
        self.message = ""

        # count prefix
        if key in range(ord("1"), ord("9") + 1) and not self.pending:
            self.repeat_count += chr(key)
            return None
        if key == ord("0") and self.repeat_count:
            self.repeat_count += "0"
            return None

        count = self.get_count()

        # pending sequences (d, y, c, g, etc.)
        if self.pending:
            return self._handle_pending(key, count)

        # exit editor (Esc or left arrow at column 0)
        if key == 27:
            result = self.prompt_save_exit()
            if result == "quit":
                return "quit"
            return None
        elif key == curses.KEY_LEFT and self.cx == 0:
            result = self.prompt_save_exit()
            if result == "quit":
                return "quit"
            return None

        # movement
        if key in (curses.KEY_LEFT, ord("h")):
            self.cx = max(0, self.cx - count)
        elif key in (curses.KEY_DOWN, ord("j")):
            self.cy = min(len(self.lines) - 1, self.cy + count)
        elif key in (curses.KEY_UP, ord("k")):
            self.cy = max(0, self.cy - count)
        elif key in (curses.KEY_RIGHT, ord("l")):
            self.cx = min(len(self.cur_line()) - 1, self.cx + count) if self.cur_line() else 0
        elif key == ord("w"):
            for _ in range(count):
                self.word_forward()
        elif key == ord("b"):
            for _ in range(count):
                self.word_backward()
        elif key == ord("e"):
            for _ in range(count):
                self.word_end()
        elif key == ord("0"):
            self.cx = 0
        elif key == ord("$"):
            self.cx = max(0, len(self.cur_line()) - 1)
        elif key == ord("^"):
            line = self.cur_line()
            self.cx = len(line) - len(line.lstrip())

        # scrolling
        elif key == curses.KEY_PPAGE or key == 2:  # ctrl-b
            h = self.stdscr.getmaxyx()[0] - 2
            self.cy = max(0, self.cy - h)
        elif key == curses.KEY_NPAGE or key == 6:  # ctrl-f
            h = self.stdscr.getmaxyx()[0] - 2
            self.cy = min(len(self.lines) - 1, self.cy + h)
        elif key == 4:  # ctrl-d
            h = (self.stdscr.getmaxyx()[0] - 2) // 2
            self.cy = min(len(self.lines) - 1, self.cy + h)
        elif key == 21:  # ctrl-u
            h = (self.stdscr.getmaxyx()[0] - 2) // 2
            self.cy = max(0, self.cy - h)

        # paragraph motions
        elif key == ord("{"):
            for _ in range(count):
                self.para_backward()
        elif key == ord("}"):
            for _ in range(count):
                self.para_forward()

        # folding
        elif key == ord("z"):
            self.pending = "z"

        # entering insert mode
        elif key == ord("i"):
            self.mode = self.INSERT
        elif key == ord("I"):
            line = self.cur_line()
            self.cx = len(line) - len(line.lstrip())
            self.mode = self.INSERT
        elif key == ord("a"):
            self.cx = min(self.cx + 1, len(self.cur_line()))
            self.mode = self.INSERT
        elif key == ord("A"):
            self.cx = len(self.cur_line())
            self.mode = self.INSERT
        elif key == ord("o"):
            self.save_undo()
            indent = self._get_indent()
            self.lines.insert(self.cy + 1, indent)
            self.cy += 1
            self.cx = len(indent)
            self.mode = self.INSERT
            self.modified = True
        elif key == ord("O"):
            self.save_undo()
            indent = self._get_indent()
            self.lines.insert(self.cy, indent)
            self.cx = len(indent)
            self.mode = self.INSERT
            self.modified = True
        elif key == ord("s"):
            self.save_undo()
            self.delete_char()
            self.mode = self.INSERT
        elif key == ord("S"):
            self.save_undo()
            indent = self._get_indent()
            self.lines[self.cy] = indent
            self.cx = len(indent)
            self.mode = self.INSERT
            self.modified = True
        elif key == ord("C"):
            self.save_undo()
            self.delete_to_eol()
            self.mode = self.INSERT

        # editing in normal mode
        elif key == ord("x"):
            self.save_undo()
            for _ in range(count):
                if self.cur_line():
                    ch = self.cur_line()[self.cx:self.cx+1]
                    self.yank_buf = [ch]
                    self.yank_linewise = False
                    self.delete_char()
        elif key == ord("X"):
            self.save_undo()
            for _ in range(count):
                if self.cx > 0:
                    self.cx -= 1
                    self.delete_char()
        elif key == ord("r"):
            self.pending = "r"
            return None
        elif key == ord("J"):
            self.save_undo()
            if self.cy < len(self.lines) - 1:
                cur = self.lines[self.cy].rstrip()
                nxt = self.lines[self.cy + 1].lstrip()
                self.lines[self.cy] = cur + " " + nxt
                self.lines.pop(self.cy + 1)
                self.cx = len(cur)
                self.modified = True
        elif key == ord("~"):
            self.save_undo()
            line = self.cur_line()
            if self.cx < len(line):
                ch = line[self.cx]
                ch = ch.upper() if ch.islower() else ch.lower()
                self.lines[self.cy] = line[:self.cx] + ch + line[self.cx+1:]
                self.cx = min(self.cx + 1, max(0, len(self.lines[self.cy]) - 1))
                self.modified = True

        # operators
        elif key == ord("d"):
            self.pending = "d"
        elif key == ord("y"):
            self.pending = "y"
        elif key == ord("c"):
            self.pending = "c"
        elif key == ord(">"):
            self.pending = ">"
        elif key == ord("<"):
            self.pending = "<"

        # jump
        elif key == ord("G"):
            if count > 1 or self.repeat_count:
                self.cy = min(count - 1, len(self.lines) - 1)
            else:
                self.cy = len(self.lines) - 1
            self.cx = 0
        elif key == ord("g"):
            self.pending = "g"

        # paste
        elif key == ord("p"):
            self.paste_after()
        elif key == ord("P"):
            self.paste_before()

        # undo/redo
        elif key == ord("u"):
            self.undo()
        elif key == 18:  # ctrl-r
            self.redo()

        # visual
        elif key == ord("v"):
            self.mode = self.VISUAL
            self.visual_start = (self.cy, self.cx)
        elif key == ord("V"):
            self.mode = self.VISUAL_LINE
            self.visual_start = (self.cy, self.cx)

        # search
        elif key == ord("/"):
            self.mode = self.SEARCH
            self.search_dir = 1
            self.search_query = ""
        elif key == ord("?"):
            self.mode = self.SEARCH
            self.search_dir = -1
            self.search_query = ""
        elif key == ord("n"):
            self.do_search(self.search_dir)
        elif key == ord("N"):
            self.do_search(-self.search_dir)
        elif key == ord("*"):
            # search word under cursor
            line = self.cur_line()
            if line:
                start = end = self.cx
                while start > 0 and (line[start-1].isalnum() or line[start-1] == "_"):
                    start -= 1
                while end < len(line) and (line[end].isalnum() or line[end] == "_"):
                    end += 1
                word = line[start:end]
                if word:
                    self.last_search = word
                    self.search_dir = 1
                    self.do_search()

        # command mode
        elif key == ord(":"):
            self.mode = self.COMMAND
            self.cmd_buf = ""

        # indent
        elif key == ord(">"):
            self.save_undo()
            self.indent_line(self.cy, 1)
        elif key == ord("<"):
            self.save_undo()
            self.indent_line(self.cy, -1)

        # dot (simplified: repeat last insert? nah, just a msg for now)
        elif key == ord("."):
            self.message = "(dot repeat not yet implemented)"

        self.clamp_cursor()
        return None

    def _handle_pending(self, key, count):
        op = self.pending
        self.pending = ""

        if op == "r" and 32 <= key <= 126:
            self.save_undo()
            line = self.cur_line()
            if self.cx < len(line):
                self.lines[self.cy] = line[:self.cx] + chr(key) + line[self.cx+1:]
                self.modified = True
            return None

        if op == "z":
            if key == ord("M"):
                self.fold_all_close()
            elif key == ord("O"):
                self.fold_all_open()
            elif key == ord("c"):
                self.fold_close(self.cy)
            elif key == ord("o"):
                self.fold_open(self.cy)
            elif key == ord("a"):
                # toggle fold at cursor
                if self.cy in self.folds:
                    self.fold_open(self.cy)
                else:
                    self.fold_close(self.cy)
            elif key == ord("z"):
                # zz: center cursor on screen
                text_h = self.stdscr.getmaxyx()[0] - 2
                vis = self.visible_lines()
                for vi, li in enumerate(vis):
                    if li == self.cy:
                        self.scroll_y = max(0, vi - text_h // 2)
                        break
            self.ensure_cursor_visible()
            self.clamp_cursor()
            return None

        if op == "g":
            if key == ord("g"):
                self.cy = 0
                self.cx = 0
            elif key == ord("d"):
                pass  # go to definition not supported
            self.clamp_cursor()
            return None

        if op == "d":
            self.save_undo()
            if key == ord("d"):
                self.delete_line(count)
            elif key == ord("w"):
                # delete word
                start = self.cx
                for _ in range(count):
                    self.word_forward()
                line = self.cur_line()
                deleted = line[start:self.cx]
                self.lines[self.cy] = line[:start] + line[self.cx:]
                self.yank_buf = [deleted]
                self.yank_linewise = False
                self.cx = start
                self.modified = True
            elif key == ord("$"):
                self.delete_to_eol()
            elif key == ord("0"):
                line = self.cur_line()
                deleted = line[:self.cx]
                self.lines[self.cy] = line[self.cx:]
                self.yank_buf = [deleted]
                self.yank_linewise = False
                self.cx = 0
                self.modified = True
            elif key == ord("G"):
                self.save_undo()
                self.yank_buf = self.lines[self.cy:]
                self.yank_linewise = True
                self.lines = self.lines[:self.cy] or [""]
                self.modified = True
            elif key == ord("g"):
                self.pending = "dg"
                return None

        elif op == "dg":
            if key == ord("g"):
                self.yank_buf = self.lines[:self.cy + 1]
                self.yank_linewise = True
                self.lines = self.lines[self.cy + 1:] or [""]
                self.cy = 0
                self.cx = 0
                self.modified = True
            return None

        elif op == "y":
            if key == ord("y"):
                self.yank_line(count)
            elif key == ord("w"):
                start = self.cx
                old_cy = self.cy
                for _ in range(count):
                    self.word_forward()
                if self.cy == old_cy:
                    self.yank_buf = [self.cur_line()[start:self.cx]]
                self.yank_linewise = False
                self.cy = old_cy
                self.cx = start
            self.clamp_cursor()
            return None

        elif op == "c":
            self.save_undo()
            if key == ord("c"):
                indent = self._get_indent()
                self.yank_buf = [self.cur_line()]
                self.yank_linewise = False
                self.lines[self.cy] = indent
                self.cx = len(indent)
                self.mode = self.INSERT
                self.modified = True
            elif key == ord("w"):
                start = self.cx
                for _ in range(count):
                    self.word_forward()
                line = self.cur_line()
                self.lines[self.cy] = line[:start] + line[self.cx:]
                self.cx = start
                self.mode = self.INSERT
                self.modified = True

        elif op == ">":
            if key == ord(">"):
                self.save_undo()
                for i in range(count):
                    if self.cy + i < len(self.lines):
                        self.indent_line(self.cy + i, 1)

        elif op == "<":
            if key == ord("<"):
                self.save_undo()
                for i in range(count):
                    if self.cy + i < len(self.lines):
                        self.indent_line(self.cy + i, -1)

        self.clamp_cursor()
        return None

    def _get_indent(self):
        line = self.cur_line()
        indent = ""
        for ch in line:
            if ch in " \t":
                indent += ch
            else:
                break
        return indent

    # ── Insert mode ──────────────────────────────────────────────────────
    def handle_insert(self, key):
        if key == 27:  # Esc
            self.save_undo()
            self.mode = self.NORMAL
            self.cx = max(0, self.cx - 1)
            self.clamp_cursor()
            return

        if key in (curses.KEY_LEFT,):
            self.cx = max(0, self.cx - 1)
        elif key in (curses.KEY_RIGHT,):
            self.cx = min(len(self.cur_line()), self.cx + 1)
        elif key in (curses.KEY_UP,):
            self.cy = max(0, self.cy - 1)
            self.cx = min(self.cx, len(self.cur_line()))
        elif key in (curses.KEY_DOWN,):
            self.cy = min(len(self.lines) - 1, self.cy + 1)
            self.cx = min(self.cx, len(self.cur_line()))
        elif key in (curses.KEY_BACKSPACE, 127, 8):
            if self.cx > 0:
                line = self.cur_line()
                self.lines[self.cy] = line[:self.cx - 1] + line[self.cx:]
                self.cx -= 1
                self.modified = True
            elif self.cy > 0:
                prev_len = len(self.lines[self.cy - 1])
                self.lines[self.cy - 1] += self.lines[self.cy]
                self.lines.pop(self.cy)
                self.cy -= 1
                self.cx = prev_len
                self.modified = True
        elif key in (curses.KEY_DC,):
            self.delete_char()
        elif key in (curses.KEY_ENTER, 10, 13):
            line = self.cur_line()
            # auto-indent
            indent = self._get_indent()
            # extra indent after colon
            stripped = line[:self.cx].rstrip()
            if stripped.endswith(":"):
                indent += "    "
            rest = line[self.cx:]
            self.lines[self.cy] = line[:self.cx]
            self.lines.insert(self.cy + 1, indent + rest)
            self.cy += 1
            self.cx = len(indent)
            self.modified = True
        elif key == 9:  # tab
            line = self.cur_line()
            self.lines[self.cy] = line[:self.cx] + "    " + line[self.cx:]
            self.cx += 4
            self.modified = True
        elif key == curses.KEY_HOME:
            self.cx = 0
        elif key == curses.KEY_END:
            self.cx = len(self.cur_line())
        elif 32 <= key <= 126 or key >= 128:
            line = self.cur_line()
            try:
                ch = chr(key)
                self.lines[self.cy] = line[:self.cx] + ch + line[self.cx:]
                self.cx += 1
                self.modified = True
            except (ValueError, OverflowError):
                pass

    # ── Visual mode ──────────────────────────────────────────────────────
    def handle_visual(self, key):
        if key == 27:  # Esc
            self.mode = self.NORMAL
            return
        # movement (same as normal)
        if key in (curses.KEY_DOWN, ord("j")):
            self.cy = min(len(self.lines) - 1, self.cy + 1)
        elif key in (curses.KEY_UP, ord("k")):
            self.cy = max(0, self.cy - 1)
        elif key in (curses.KEY_LEFT, ord("h")):
            self.cx = max(0, self.cx - 1)
        elif key in (curses.KEY_RIGHT, ord("l")):
            self.cx = min(max(0, len(self.cur_line()) - 1), self.cx + 1)
        elif key == ord("w"):
            self.word_forward()
        elif key == ord("b"):
            self.word_backward()
        elif key == ord("e"):
            self.word_end()
        elif key == ord("$"):
            self.cx = max(0, len(self.cur_line()) - 1)
        elif key == ord("0"):
            self.cx = 0
        elif key == ord("G"):
            self.cy = len(self.lines) - 1
        elif key == ord("g"):
            self.pending = "g"
            # simplified: just go to top
            self.cy = 0
            self.cx = 0

        # operations
        elif key == ord("d") or key == ord("x"):
            self.save_undo()
            self.visual_delete()
        elif key == ord("y"):
            self.visual_yank()
        elif key == ord("c"):
            self.save_undo()
            self.visual_delete()
            self.mode = self.INSERT
        elif key == ord(">"):
            self.save_undo()
            sy = min(self.visual_start[0], self.cy)
            ey = max(self.visual_start[0], self.cy)
            for i in range(sy, ey + 1):
                self.indent_line(i, 1)
            self.mode = self.NORMAL
        elif key == ord("<"):
            self.save_undo()
            sy = min(self.visual_start[0], self.cy)
            ey = max(self.visual_start[0], self.cy)
            for i in range(sy, ey + 1):
                self.indent_line(i, -1)
            self.mode = self.NORMAL
        elif key == ord(":"):
            self.mode = self.COMMAND
            self.cmd_buf = ""

        self.clamp_cursor()

    # ── Save prompt ───────────────────────────────────────────────────────
    def prompt_save_exit(self):
        """Prompt to save. Enter=save, n=don't save, Esc=cancel."""
        if not self.modified:
            return "quit"
        h, w = self.stdscr.getmaxyx()
        msg = " Save changes? [Enter]=save  [n]=don't save  [Esc]=cancel "
        try:
            self.stdscr.addnstr(h - 1, 0, msg.ljust(w), w, self.cp("statusbar") | curses.A_BOLD)
        except curses.error:
            pass
        self.stdscr.refresh()
        while True:
            k = self.stdscr.getch()
            if k in (curses.KEY_ENTER, 10, 13):
                self.save_file()
                return "quit"
            elif k == ord("n"):
                return "quit"
            elif k == 27:
                return None  # cancel, stay in editor

    # ── Main loop ────────────────────────────────────────────────────────
    def run(self):
        while True:
            self.draw()
            key = self.stdscr.getch()

            if key == curses.KEY_RESIZE:
                continue

            if self.mode == self.NORMAL:
                result = self.handle_normal(key)
                if result == "quit":
                    return
            elif self.mode == self.INSERT:
                self.handle_insert(key)
            elif self.mode in (self.VISUAL, self.VISUAL_LINE):
                self.handle_visual(key)
            elif self.mode == self.COMMAND:
                if key == 27:
                    self.mode = self.NORMAL
                elif key in (curses.KEY_ENTER, 10, 13):
                    result = self.exec_command()
                    if result == "quit":
                        return
                elif key in (curses.KEY_BACKSPACE, 127, 8):
                    if self.cmd_buf:
                        self.cmd_buf = self.cmd_buf[:-1]
                    else:
                        self.mode = self.NORMAL
                elif 32 <= key <= 126:
                    self.cmd_buf += chr(key)
            elif self.mode == self.SEARCH:
                if key == 27:
                    self.mode = self.NORMAL
                elif key in (curses.KEY_ENTER, 10, 13):
                    self.last_search = self.search_query
                    self.mode = self.NORMAL
                    self.do_search()
                elif key in (curses.KEY_BACKSPACE, 127, 8):
                    if self.search_query:
                        self.search_query = self.search_query[:-1]
                    else:
                        self.mode = self.NORMAL
                elif 32 <= key <= 126:
                    self.search_query += chr(key)


def main(stdscr):
    filename = sys.argv[1] if len(sys.argv) > 1 else None
    Editor(stdscr, filename).run()


if __name__ == "__main__":
    curses.wrapper(main)
