"""Keyboard-driven tree picker used by choose_tab.py."""

import json
import re
import sys
import unicodedata

from kittens.tui.handler import Handler, result_handler
from kittens.tui.loop import Loop
from kitty.fast_data_types import truncate_point_for_length, wcswidth


RESET = '\x1b[0m'
CONTROLS = re.compile(
    r'\x1b(?:\][^\x07\x1b]*(?:\x07|\x1b\\)|[PX^_].*?\x1b\\|'
    r'\[[0-?]*[ -/]*[@-~]|[@-_])', re.DOTALL)
SGR = re.compile(r'\x1b\[[0-9;:]*m\Z')


def clean_text(text):
    text = CONTROLS.sub('', text)
    return ''.join(c for c in text if not unicodedata.category(c).startswith('C'))


def clip_text(text, width):
    end = truncate_point_for_length(text, width)
    while end < len(text) and unicodedata.category(text[end]) in ('Mn', 'Me'):
        end += 1
    return text[:end]


def clip_ansi(text, width):
    """Retain only SGR formatting; never let a snapshot control the terminal."""
    result = []
    used = 0
    end = 0
    for match in CONTROLS.finditer(text):
        plain = clean_text(text[end:match.start()])
        piece = clip_text(plain, max(0, width - used))
        result.append(piece)
        used += wcswidth(piece)
        if len(piece) != len(plain):
            return ''.join(result) + RESET
        if SGR.fullmatch(match.group()):
            result.append(match.group())
        end = match.end()
    plain = clean_text(text[end:])
    result.append(clip_text(plain, max(0, width - used)))
    return ''.join(result) + RESET


class Preview:
    """Position a colored, cell-clipped preview inside a terminal rectangle."""

    def __init__(self, top, width, height):
        self.top, self.width, self.height = top, width, height
        self.output = []

    def put(self, x, y, text, width):
        if 0 <= x < self.width and 0 <= y < self.height:
            width = min(width, self.width - x)
            if width > 0:
                self.output.append(f'\x1b[{self.top + y};{x + 1}H{RESET}'
                                   + clip_ansi(text, width))

    def box(self, x, y, width, height, title='', active=False):
        if width < 2 or height < 2:
            return
        color = '\x1b[33m' if active else '\x1b[90m'
        self.put(x, y, color + '┌' + '─' * (width - 2) + '┐', width)
        for row in range(y + 1, y + height - 1):
            self.put(x, row, color + '│', 1)
            self.put(x + width - 1, row, color + '│', 1)
        self.put(x, y + height - 1, color + '└' + '─' * (width - 2) + '┘', width)
        if title:
            self.put(x + 1, y, color + ' ' + clean_text(title) + ' ', width - 2)

    def badge(self, x, y, width, height, label, active):
        if width < 5 or height < 3:
            return
        label = clip_text(clean_text(label), width - 4)
        badge_width = wcswidth(label) + 4
        bx = x + (width - badge_width) // 2
        by = y + (height - 3) // 2
        self.box(bx, by, badge_width, 3)
        color = '\x1b[33m' if active else '\x1b[34m'
        self.put(bx + 1, by + 1, color + ' ' + label + ' ', badge_width - 2)

    def tab(self, snapshot, x, y, width, height, title=None, active=False):
        panes = snapshot.get('panes', []) if snapshot else []
        if not panes:
            self.put(x, y, 'Preview unavailable', width)
            if title is not None:
                self.badge(x, y, width, height, title, active)
            return
        left = min(p['rect'][0] for p in panes)
        top = min(p['rect'][1] for p in panes)
        right = max(p['rect'][2] for p in panes)
        bottom = max(p['rect'][3] for p in panes)
        small = 0
        dividers = []
        badges = []
        for index, pane in enumerate(panes, 1):
            l, t, r, b = pane['rect']
            px = x + round((l - left) * width / max(1, right - left))
            py = y + round((t - top) * height / max(1, bottom - top))
            pw = x + round((r - left) * width / max(1, right - left)) - px
            ph = y + round((b - top) * height / max(1, bottom - top)) - py
            if l > left:
                dividers.extend((px, row, '│', 1) for row in range(py, py + ph))
                px, pw = px + 1, pw - 1
            if t > top:
                dividers.append((px, py, '─' * max(0, pw), pw))
                py, ph = py + 1, ph - 1
            if pw < 1 or ph < 1:
                small += 1
                continue
            lines = pane.get('lines')
            for row, line in enumerate((lines if lines is not None else ['Preview unavailable'])[:ph]):
                self.put(px, py + row, line, pw)
            if title is None:
                badges.append((px, py, pw, ph, str(index), pane['active']))
        for dx, dy, text, count in dividers:
            self.put(dx, dy, '\x1b[90m' + text, count)
        for badge in badges:
            self.badge(*badge)
        if title is not None:
            self.badge(x, y, width, height, title, active)
        if small:
            self.put(x, y + height - 1, f'{small} pane(s) too small', width)


class WindowPicker(Handler):
    def __init__(self, entries, collapsed=False, snapshots=None):
        self.all_entries = entries
        self.groups = {i for i, entry in enumerate(entries) if entry['kind'] == 'window'}
        self.collapsed = set(self.groups) if collapsed else set()
        self.selected = 0
        self.offset = 0
        self.answer = None
        self.query = ''
        self.searching = False
        self.search_origin = ('', 0)
        self.snapshots = snapshots or {}
        self.preview_pages = {}

    def visible_indices(self):
        hidden = False
        visible = []
        for index in range(len(self.all_entries)):
            if index in self.groups:
                hidden = index in self.collapsed
                visible.append(index)
            elif not hidden:
                visible.append(index)
        return visible

    @property
    def entries(self):
        return [self.all_entries[i] for i in self.visible_indices()]

    def expand_group(self, expand):
        visible = self.visible_indices()
        if not visible:
            return
        index = visible[self.selected]
        if index not in self.groups:
            return
        if expand:
            self.collapsed.discard(index)
        else:
            self.collapsed.add(index)
        self.draw()

    def matches(self):
        if not self.query:
            return []
        query = self.query.casefold()
        return [i for i, entry in enumerate(self.entries)
                if query in clean_text(entry['label']).casefold()]

    def update_search(self):
        matches = self.matches()
        if matches and self.selected not in matches:
            self.selected = next((i for i in matches if i >= self.selected), matches[0])
        self.draw()

    def next_match(self, delta):
        matches = self.matches()
        if matches:
            ordered = matches if delta > 0 else list(reversed(matches))
            self.selected = next(
                (i for i in ordered if (i - self.selected) * delta > 0), ordered[0])
        self.draw()

    def initialize(self):
        self.draw()

    def preview_height(self):
        if self.screen_size.rows < 18 or self.screen_size.cols < 24:
            return 0
        return (self.screen_size.rows - 1) // 2

    def page_preview(self, delta):
        if self.entries and self.preview_height():
            entry = self.entries[self.selected]
            if entry['kind'] == 'window':
                capacity = max(1, (self.screen_size.cols - 2) // 24)
                count = len(entry['tabs'])
                default = entry['tabs'].index(entry['tab_id']) // capacity
                page = self.preview_pages.get(id(entry), default)
                self.preview_pages[id(entry)] = max(0, min((count - 1) // capacity, page + delta))
                self.draw()

    def draw_preview(self, height):
        preview = Preview(self.screen_size.rows - height, self.screen_size.cols, height)
        if not self.entries:
            preview.box(0, 0, preview.width, height, 'Snapshot')
            return ''.join(preview.output)
        entry = self.entries[self.selected]
        tab_ids = entry.get('tabs', [entry['tab_id']])
        capacity = max(1, (preview.width - 2) // 24) if entry['kind'] == 'window' else 1
        pages = max(1, (len(tab_ids) + capacity - 1) // capacity)
        default = tab_ids.index(entry['tab_id']) // capacity
        page = min(pages - 1, self.preview_pages.get(id(entry), default))
        self.preview_pages[id(entry)] = page
        start = page * capacity
        shown = tab_ids[start:start + capacity]
        label = f'{entry["title"]} — snapshot'
        if pages > 1:
            label += f' — tabs {start + 1}–{start + len(shown)}/{len(tab_ids)} [{page + 1}/{pages}]'
        preview.box(0, 0, preview.width, height, label)
        inner_width = preview.width - 2
        for offset, tab_id in enumerate(shown):
            left = 1 + offset * inner_width // len(shown)
            right = 1 + (offset + 1) * inner_width // len(shown)
            if offset:
                for row in range(1, height - 1):
                    preview.put(left, row, '\x1b[90m│', 1)
                left += 1
            snapshot = self.snapshots.get(str(tab_id))
            title = None
            if entry['kind'] == 'window':
                title = f'{start + offset + 1}:{snapshot["title"] if snapshot else "Tab"}'
            active = tab_id == entry['tab_id']
            preview.tab(snapshot, left, 1, right - left, height - 2,
                        title, active)
        return ''.join(preview.output)

    def draw(self):
        height = self.preview_height()
        rows = max(1, self.screen_size.rows - 1 - height)
        cols = max(1, self.screen_size.cols)
        self.offset = min(self.offset, self.selected)
        self.offset = max(self.offset, self.selected - rows + 1)
        output = ["\x1b[?25l\x1b[H\x1b[2J"]
        matches = self.matches()
        visible = self.visible_indices()
        for index in range(self.offset, min(len(self.entries), self.offset + rows)):
            label = self.entries[index]['label']
            if visible[index] in self.collapsed:
                label = '+' + label[1:]
            # Titles are untrusted terminal text: strip control sequences/characters.
            label = clean_text(label)
            line = f"({index}) {label}"
            line = line[:truncate_point_for_length(line, cols)]
            style = "\x1b[1;33m" if index in matches else ""
            if index == self.selected:
                style += "\x1b[7m"
            output.append(f"{style}{line}\x1b[K\x1b[0m\r\n")
        if height:
            output.append(self.draw_preview(height))
        footer = 'j/k ↑/↓ move  h/l ←/→ fold  / search'
        footer += '  [] page' if height else ''
        footer += '  Enter select  Esc cancel'
        if not height:
            footer += '  Preview needs 24 cols × 18 rows'
        if self.searching or self.query:
            position = str(matches.index(self.selected) + 1) if self.selected in matches else '-'
            status = f"{position}/{len(matches)}" if matches else "no matches"
            help_text = "Enter: finish search  Esc: undo" if self.searching else "n/N: next/previous  /: search"
            footer = f"/{self.query} [{status}]  {help_text}"
        footer = footer[:truncate_point_for_length(footer, cols)]
        output.append(f"\x1b[0m\x1b[{self.screen_size.rows};1H{footer}")
        self.write(''.join(output))

    def move(self, delta):
        self.selected = max(0, min(len(self.entries) - 1, self.selected + delta))
        self.draw()

    def on_text(self, text, in_bracketed_paste=False):
        if len(text) > 1 and not in_bracketed_paste:
            for char in text:
                self.on_text(char)
            return
        if self.searching:
            self.query += ''.join(c for c in text if not unicodedata.category(c).startswith('C'))
            self.update_search()
        elif not in_bracketed_paste:
            if text == '/':
                self.search_origin = (self.query, self.selected)
                self.query = ''
                self.searching = True
                self.draw()
            elif text == 'n':
                self.next_match(1)
            elif text == 'N':
                self.next_match(-1)
            elif text == 'j':
                self.move(1)
            elif text == 'k':
                self.move(-1)
            elif text == 'h':
                self.expand_group(False)
            elif text == 'l':
                self.expand_group(True)
            elif text == '[':
                self.page_preview(-1)
            elif text == ']':
                self.page_preview(1)
            elif text == 'q':
                self.quit_loop(0)

    def on_key(self, key):
        if key.is_release:
            return
        if self.searching:
            if key.matches('enter'):
                self.searching = False
                self.draw()
                return
            if key.matches('escape'):
                self.query, self.selected = self.search_origin
                self.searching = False
                self.draw()
                return
            if key.matches('backspace'):
                self.query = self.query[:-1]
                self.update_search()
                return
            if key.matches('ctrl+u'):
                self.query = ''
                self.update_search()
                return
            if key.matches('left') or key.matches('right'):
                return
        if key.matches('down'):
            self.move(1)
        elif key.matches('up'):
            self.move(-1)
        elif key.matches('left'):
            self.expand_group(False)
        elif key.matches('right'):
            self.expand_group(True)
        elif key.matches('enter'):
            if self.entries:
                self.answer = self.entries[self.selected]['tab_id']
            self.quit_loop(0)
        elif key.matches('escape'):
            self.quit_loop(0)

    def on_resize(self, size):
        super().on_resize(size)
        self.draw()

    def on_interrupt(self):
        self.quit_loop(0)

    def on_eot(self):
        self.quit_loop(0)


def main(args):
    payload = json.load(sys.stdin)
    picker = WindowPicker(payload['entries'], collapsed='--collapsed' in args[1:],
                          snapshots=payload['snapshots'])
    Loop().loop(picker)
    return picker.answer


@result_handler()
def handle_result(args, answer, target_window_id, boss):
    # choose_tab.py supplies the callback and activates after this overlay closes.
    pass