#!/usr/bin/env python3
"""
Kitty kitten: select a tab from a list with fuzzy search.

Usage: map some_key kitten select_tab.py

Keys:
  Up/Down, Page Up/Page Down - scroll list
  Enter - select highlighted tab
  Any other input - search box (incremental fuzzy filter)
  Ctrl+Up / Ctrl+Down - search history (previous/next)
  Ctrl+C / Escape - quit without selecting
"""

from __future__ import annotations

# from kitty.typing_compat import KeyEventType
from kitty.key_encoding import EventType
from kitty.boss import Boss
from kittens.tui.handler import Handler, kitten_ui
from kittens.tui.loop import Loop, debug
from kittens.tui.line_edit import LineEdit
from kittens.tui.operations import styled

import json
import os
import sys
import pprint
from typing import Any


def format_cwd_like_zsh(path: str) -> str:
    """Format CWD like zsh %(4~|…/%3~|%~): full path, or …/last/3/parts if 4+ segments."""
    if not path:
        return ""
    home = os.environ.get("HOME", "") or os.path.expanduser("~")
    if home and (path == home or path.startswith(home + "/")):
        path = "~" + path[len(home) :].rstrip("/") or ""
    path = path.rstrip("/") or "/"
    segments = [s for s in path.split("/") if s]
    if len(segments) >= 4:
        return "…/" + "/".join(segments[-3:])
    return path if path else "~"


class TabRow:
    """One row of data for a tab (for display and filtering)."""

    def __init__(
        self,
        tab_id,
        os_window_id,
        index_in_os_window,
        title,
        cwd,
        vt,
        foreground,
        search_text="",
    ):
        self.tab_id = tab_id
        self.os_window_id = os_window_id
        self.index_in_os_window = index_in_os_window
        self.title = title
        self.cwd = cwd
        self.vt = vt
        self.foreground = foreground
        parts = [title, cwd]
        if vt is not None:
            parts.append(vt)
        parts.append(foreground)
        self.search_text = " ".join(parts).lower() if not search_text else search_text


def fuzzy_match(query: str, text: str) -> bool:
    """True if query characters appear in text in order (case-insensitive)."""
    if not query:
        return True
    q = query.lower()
    t = text.lower()
    i = 0
    for c in q:
        j = t.find(c, i)
        if j < 0:
            return False
        i = j + 1
    return True


def get_fuzzy_match_indices(query: str, text: str) -> list[int]:
    """Indices in text that match each query character in order (case-insensitive)."""
    if not query:
        return []
    q = query.lower()
    t = text.lower()
    indices = []
    i = 0
    for c in q:
        j = t.find(c, i)
        if j < 0:
            break
        indices.append(j)
        i = j + 1
    return indices


def highlight_matches(
    segment_str: str,
    match_indices: list[int],
    content_len: int,
) -> str:
    """Wrap matching character runs in segment_str with styled bold cyan (or fallback SGR)."""
    in_content = {i for i in match_indices if 0 <= i < content_len}
    if not in_content:
        return segment_str
    # Merge consecutive indices into (start, end) runs
    sorted_i = sorted(in_content)
    runs = []
    start = sorted_i[0]
    end = start + 1
    for i in sorted_i[1:]:
        if i == end:
            end = i + 1
        else:
            runs.append((start, end))
            start, end = i, i + 1
    runs.append((start, end))
    # Build string with highlights
    out = []
    pos = 0
    hi_on = "\033[1;36m"  # bold cyan
    hi_off = "\033[22;39m"
    for s, e in runs:
        out.append(segment_str[pos:s])
        chunk = segment_str[s:e]
        st = styled(chunk, fg="cyan", bold=True)
        out.append(st if st else hi_on + chunk + hi_off)
        pos = e
    out.append(segment_str[pos:])
    return "".join(out)


def longest_substring_match_length(query: str, text: str) -> int:
    """Length of longest contiguous substring of query that appears in text (case-insensitive)."""
    if not query:
        return 0
    q = query.lower()
    t = text.lower()
    best = 0
    for start in range(len(q)):
        for end in range(start + 1, len(q) + 1):
            if q[start:end] in t:
                best = max(best, end - start)
    return best


def gather_tabs_from_ls(ls_data: list[dict[str, Any]]) -> list[TabRow]:
    """Build a flat list of TabRow from kitten @ ls JSON (list of OS windows)."""
    rows: list[TabRow] = []
    # Exclude the tab that contains the window we're running in (the overlay's parent)
    try:
        self_window_id = int(os.environ.get("KITTY_WINDOW_ID", "") or 0)
    except (ValueError, TypeError):
        self_window_id = 0
    for os_win in ls_data:
        os_window_id = os_win.get("id", 0)
        tabs = os_win.get("tabs") or []
        for tab_idx, tab in enumerate(tabs):
            windows = tab.get("windows") or []
            if self_window_id and any(w.get("id") == self_window_id for w in windows):
                continue
            tab_id = tab.get("id", 0)
            title = tab.get("title") or ""
            # Use active/focused window in tab; fallback to first
            active = next(
                (w for w in windows if w.get("is_focused") or w.get("is_active")),
                windows[0] if windows else None,
            )
            cwd = ""
            vt = None
            foreground = ""
            if active:
                debug(f"Active window in tab {tab_id}: {active}")
                cwd = active.get("cwd") or ""
                user_vars = active.get("user_vars")
                if isinstance(user_vars, dict):
                    vt = user_vars.get("vt")
                else:
                    vt = None
                # Foreground process: cmdline or foreground_processes
                cmdline = active.get("cmdline")
                if cmdline:
                    foreground = (
                        " ".join(cmdline) if isinstance(cmdline, list) else str(cmdline)
                    )
                else:
                    fg_procs = active.get("foreground_processes")
                    if fg_procs:
                        parts = []
                        for p in fg_procs:
                            if isinstance(p, dict):
                                pcmd = p.get("cmdline")
                                if pcmd:
                                    parts.append(
                                        " ".join(pcmd)
                                        if isinstance(pcmd, list)
                                        else str(pcmd)
                                    )
                            else:
                                parts.append(str(p))
                        foreground = " | ".join(parts)
            rows.append(
                TabRow(
                    tab_id=tab_id,
                    os_window_id=os_window_id,
                    index_in_os_window=tab_idx,
                    title=title or "(no title)",
                    cwd=cwd,
                    vt=vt,
                    foreground=foreground or "(no process)",
                    search_text="",
                )
            )
    return rows


class TabSelectHandler(Handler):
    """TUI handler: list of tabs, search box, scroll, history."""

    def __init__(self, tab_rows: list[TabRow]) -> None:
        super().__init__()
        self.all_rows = tab_rows
        self.filtered_rows = tab_rows[:]
        # self.search_history: list[str] = []
        # self.history_index = -1  # -1 = current (not in history)
        self.selected_index = 0  # index in filtered list
        self.scroll_offset = 0
        self.selected_tab_result: dict[str, Any] | None = None  # set on Enter
        self.line_edit = LineEdit()

    # Minimum size so we always draw something (overlay PTY can report 0x0 at start)
    MIN_ROWS = 24
    MIN_COLS = 80

    @property
    def screen_rows(self) -> int:
        return max(self.MIN_ROWS, getattr(self.screen_size, "rows", 0) or 0)

    @property
    def screen_cols(self) -> int:
        return max(self.MIN_COLS, getattr(self.screen_size, "cols", 0) or 0)

    def initialize(self) -> None:
        # self.print(f"In TabSelectHandler...")
        self.cmd.clear_screen()
        self.cmd.set_cursor_visible(True)
        self._apply_filter()
        self._clamp_selection()
        # Draw immediately and also schedule a draw on next event-loop tick. When run
        # as an overlay, the PTY size can be 0x0 at first or kitty may only show
        # content sent after overlay_ready(); the delayed draw ensures something
        # appears once the overlay is ready.
        self._draw()
        try:
            self.asyncio_loop.call_soon(self._draw)
        except Exception:
            pass
        # self.write(SAVE_CURSOR)
        # debug("Initialization complete")

    def _apply_filter(self) -> None:
        # debug("In _apply_filter")
        q = self.line_edit.current_input.strip()
        if not q:
            self.filtered_rows = self.all_rows[:]
        else:
            q_lower = q.lower()
            # Exact VT match: if the search term equals a tab's VT, show only that/those.
            vt_exact = [
                r for r in self.all_rows if r.vt is not None and str(r.vt).strip() == q
            ]
            if vt_exact:
                self.filtered_rows = vt_exact
            else:
                self.filtered_rows = [
                    r for r in self.all_rows if fuzzy_match(q_lower, r.search_text)
                ]
                # Sort by longest contiguous substring match (best first), then original order
                self.filtered_rows.sort(
                    key=lambda r: (
                        -longest_substring_match_length(q_lower, r.search_text),
                        self.all_rows.index(r),
                    )
                )
        self._clamp_selection()

    def on_resize(self, screen_size: Any) -> None:
        # debug(f"In on_resize with screen_size: {screen_size}")
        super().on_resize(screen_size)
        self._clamp_selection()
        self._draw()

    def _clamp_selection(self) -> None:
        # debug("In _clamp_selection")
        n = len(self.filtered_rows)
        if n == 0:
            self.selected_index = 0
            self.scroll_offset = 0
            return
        if self.selected_index >= n:
            self.selected_index = 0
        elif self.selected_index < 0:
            self.selected_index = n - 1
        # self.selected_index = max(0, min(self.selected_index, n - 1))
        # Keep selection in view
        list_top = 1
        list_bottom = self.screen_rows - 2
        list_height = list_bottom - list_top + 1
        if self.selected_index < self.scroll_offset:
            self.scroll_offset = self.selected_index
        elif self.selected_index >= self.scroll_offset + list_height:
            self.scroll_offset = self.selected_index - list_height + 1

    # def _push_search_to_history(self) -> None:
    #     debug("In _push_search_to_history")
    #     if self.search_query.strip() and (
    #         not self.search_history
    #         or self.search_history[-1] != self.search_query.strip()
    #     ):
    #         self.search_history.append(self.search_query.strip())
    #     self.history_index = -1

    def _draw(self) -> None:
        # debug("In _draw")
        self.cmd.clear_screen()
        rows, cols = self.screen_rows, self.screen_cols
        list_top = 1
        list_bottom = rows - 2
        list_height = list_bottom - list_top + 1

        # debug(f"in draw: rows: {rows}, cols: {cols}")

        # Column widths (visible chars only). Four segments, three spaces between.
        w_vt = 2
        remaining = cols - w_vt - 3
        w_title = max(12, int(remaining * 0.45))
        w_foreground = max(10, int(remaining * 0.35))
        w_cwd = max(8, remaining - w_title - w_foreground)

        # Header: plain string of exactly cols visible characters, then style once
        header_plain = (
            "VT".ljust(w_vt)
            + " "
            + "Title".ljust(w_title)
            + " "
            + "Process".ljust(w_foreground)
            + " "
            + "CWD".ljust(w_cwd)
        )
        self.cmd.set_cursor_position(0, 0)
        self.write(styled(header_plain, bold=True) + "\r\n")

        # List
        q = self.line_edit.current_input.strip().lower()
        visible = self.filtered_rows[
            self.scroll_offset : self.scroll_offset + list_height
        ]

        def segment_plain(content: str, width: int, add_dots: bool) -> tuple[str, int]:
            """Return (padded segment string, content_len for highlighting)."""
            if len(content) > width:
                if add_dots and width >= 2:
                    content = content[: width - 2] + ".."
                else:
                    content = content[:width]
            content_len = len(content)
            return content.ljust(width), content_len

        for i, row in enumerate(visible):
            y = list_top + i
            global_idx = self.scroll_offset + i
            is_selected = global_idx == self.selected_index

            # Build each segment as plain text with exact visible width (no SGR yet)
            vt_plain, vt_content_len = segment_plain(
                row.vt or "-", w_vt, add_dots=False
            )
            title_plain, title_content_len = segment_plain(
                row.title, w_title, add_dots=True
            )
            fg_plain, fg_content_len = segment_plain(
                row.foreground, w_foreground, add_dots=True
            )
            cwd_display = format_cwd_like_zsh(row.cwd)
            cwd_plain, cwd_content_len = segment_plain(
                cwd_display, w_cwd, add_dots=True
            )

            # Apply match highlighting (inserts SGR; visible length unchanged)
            if q:
                vt_styled = highlight_matches(
                    vt_plain,
                    get_fuzzy_match_indices(q, row.vt or ""),
                    vt_content_len,
                )
                title_styled = highlight_matches(
                    title_plain,
                    get_fuzzy_match_indices(q, row.title),
                    title_content_len,
                )
                fg_styled = highlight_matches(
                    fg_plain,
                    get_fuzzy_match_indices(q, row.foreground),
                    fg_content_len,
                )
                cwd_styled = highlight_matches(
                    cwd_plain,
                    get_fuzzy_match_indices(q, cwd_display),
                    cwd_content_len,
                )
            else:
                vt_styled = vt_plain
                title_styled = title_plain
                fg_styled = fg_plain
                cwd_styled = cwd_plain

            # Concatenate segments with single spaces; total visible length is fixed
            line_plain = (
                vt_styled + " " + title_styled + " " + fg_styled + " " + cwd_styled
            )

            # If selected, wrap entire line in reverse (no slicing of output)
            if is_selected:
                rev = styled(line_plain, reverse=True)
                line_out = rev if rev else "\033[7m" + line_plain + "\033[27m"
            else:
                line_out = line_plain

            self.cmd.set_cursor_position(0, y)
            self.write(line_out)
            self.cmd.clear_to_eol()
            self.write("\r\n")

        # Search line
        self.cmd.set_cursor_position(0, rows - 1)
        self.line_edit.write(
            self.write, prompt="Filter: ", screen_cols=self.screen_size.cols
        )
        # prompt = "Filter: "
        # self.write(prompt + self.search_query)
        # self.cmd.clear_to_eol()
        # self.cmd.set_cursor_position(len(prompt) + self.search_cursor, rows - 1)
        # debug(
        #     "At end of draw, flushing, current_input: {self.line_edit.current_input}, cursor_pos: {self.line_edit.cursor_pos}"
        # )
        self.flush()

    # def on_key_event(
    #     self, key_event: KeyEventType, in_bracketed_paste: bool = False
    # ) -> None:
    #     debug(f"Got key event: {key_event}")
    #     if key_event.text:
    #         self._on_search_input(key_event.text)
    #         return
    #     if self.perform_default_key_action(key_event):
    #         return
    #     self.on_key(key_event)
    #
    def on_key(self, key_event) -> None:
        # debug(f"In on_key with key_event: {key_event}")
        if self.line_edit.on_key(key_event):
            # debug(f"Handled by line_edit: {self.line_edit}")
            self.selected_index = 0
            self._clamp_selection()
            self._apply_filter()
            self._draw()
            return

        if key_event.type is EventType.RELEASE:
            return
        match key_event.key:
            case "ESCAPE":
                self.on_interrupt()
                return
            case "UP":
                self.selected_index = self.selected_index - 1
                self._clamp_selection()
                self._draw()
                return
            case "DOWN":
                self.selected_index = self.selected_index + 1
                self._clamp_selection()
                self._draw()
                return
            case "PAGE_UP":
                list_height = self.screen_rows - 3
                self.selected_index = max(0, self.selected_index - list_height)
                self._clamp_selection()
                self._draw()
                return
            case "PAGE_DOWN":
                list_height = self.screen_rows - 3
                self.selected_index = min(
                    len(self.filtered_rows) - 1, self.selected_index + list_height
                )
                self._clamp_selection()
                self._draw()
                return
            case "ENTER":
                self._select_current()
                return
        # if key_event.matches("ctrl+up"):
        #     self._history_prev()
        #     return
        # if key_event.matches("ctrl+down"):
        #     self._history_next()
        #     return
        # Line editing in search box

    def on_text(self, text: str, in_bracketed_paste: bool = False) -> None:
        # debug(f"In on_text with text: {text}, in_bracketed_paste: {in_bracketed_paste}")
        self.line_edit.on_text(text, in_bracketed_paste)
        # debug(
        #     f"After line_edit.on_text, current_input: {self.line_edit.current_input}, cursor_pos: {self.line_edit.cursor_pos}"
        # )
        self.selected_index = 0
        self._clamp_selection()
        self._apply_filter()
        self._draw()

    # def _history_prev(self) -> None:
    #     debug("In _history_prev")
    #     if not self.search_history:
    #         return
    #     if self.history_index <= 0:
    #         if self.search_query.strip():
    #             self.search_history.append(self.search_query.strip())
    #         self.history_index = len(self.search_history) - 1
    #     else:
    #         self.history_index -= 1
    #     self.search_query = self.search_history[self.history_index]
    #     self.search_cursor = len(self.search_query)
    #     self._apply_filter()
    #     self._draw()
    #
    # def _history_next(self) -> None:
    #     debug("In _history_next")
    #     if not self.search_history or self.history_index < 0:
    #         return
    #     self.history_index += 1
    #     if self.history_index >= len(self.search_history):
    #         self.history_index = -1
    #         self.search_query = ""
    #     else:
    #         self.search_query = self.search_history[self.history_index]
    #     self.search_cursor = len(self.search_query)
    #     self._apply_filter()
    #     self._draw()
    #
    def _select_current(self) -> None:
        # debug("In _select_current")
        if not self.filtered_rows:
            return
        row = self.filtered_rows[self.selected_index]
        self.selected_tab_result = {
            "tab_id": row.tab_id,
            "os_window_id": row.os_window_id,
            "index": row.index_in_os_window,
        }
        self.quit_loop(0)

    def on_interrupt(self) -> None:
        debug("In on_interrupt")
        self.quit_loop(1)

    on_eot = on_interrupt


@kitten_ui(allow_remote_control=True)
def main(args: list[str]) -> str:
    debug("In select_tab.py main")
    cp = main.remote_control(["ls"], capture_output=True)
    if cp.returncode != 0:
        debug(f"Remote control error {cp.returncode}: {cp.stderr}")
        sys.stderr.buffer.write(cp.stderr or b"")
        raise SystemExit(cp.returncode)
    ls_data = json.loads(cp.stdout.decode())
    tab_rows = gather_tabs_from_ls(ls_data)
    if not tab_rows:
        debug("No tab rows?")
        raise SystemExit(1)
    debug(f"Have tab rows {pprint.pformat(tab_rows)}")
    handler = TabSelectHandler(tab_rows)
    debug(f"Created handler: {handler}")
    sys.stdout.flush()
    loop = Loop()
    loop.loop(handler)
    debug(f"Loop exited, selected_tab_result: {handler.selected_tab_result}")

    return json.dumps(handler.selected_tab_result)


# @result_handler(has_ready_notification=True)
def handle_result(
    args: list[str],
    input: str,
    target_window_id: int,
    boss: Boss,
) -> None:
    data = json.loads(input) if input else None
    print("In handle_result with data:", data)
    if not data:
        return
    tab_id = data.get("tab_id")
    os_window_id = data.get("os_window_id")
    if tab_id is None:
        return
    # Use Boss API: match tab by id, focus its OS window if needed, then make it active
    tabs = list(boss.match_tabs(f"id:{tab_id}"))
    if not tabs:
        return
    tab = tabs[0]
    # If the tab is in a different OS window, focus that window first
    if os_window_id is not None and hasattr(boss, "focus_os_window"):
        try:
            boss.focus_os_window(os_window_id)
        except Exception:
            pass
    tab.make_active()