"""Choose a tab grouped by native Kitty window, and focus its window."""

import json

from kittens.tui.handler import result_handler
from kitty.fast_data_types import get_os_window_title


def main(args):
    pass


def snapshot_tab(tab, boss):
    """Capture before opening the chooser, without focusing or relaying out tabs."""
    panes = []
    active = tab.active_window
    try:
        groups = list(tab.windows.iter_all_layoutable_groups(only_visible=True))
    except Exception:
        groups = []
    for group in groups:
        window = boss.window_id_map.get(group.active_window_id)
        if window is None:
            continue
        try:
            geometry = window.geometry
            pane = dict(id=window.id, title=window.title,
                        active=window is active,
                        rect=[geometry.left, geometry.top, geometry.right, geometry.bottom],
                        lines=None)
        except Exception:
            continue
        try:
            pane['lines'] = [window.screen.visual_line(i).as_ansi()
                             for i in range(window.screen.lines)]
        except Exception:
            pass
        panes.append(pane)
    return dict(title=tab.name or tab.title, panes=panes)


@result_handler(no_ui=True)
def handle_result(args, answer, target_window_id, boss):
    source = boss.window_id_map.get(target_window_id)
    current = source.tabref() if source is not None else boss.active_tab

    def chosen(tab_id):
        if not isinstance(tab_id, int):
            return
        for tab in boss.all_tabs:
            if tab.id == tab_id:
                window = tab.active_window
                if window is not None:
                    boss.set_active_window(window, switch_os_window_if_needed=True)
                return

    entries = []
    snapshots = {}
    for index, (os_window_id, manager) in enumerate(boss.os_window_map.items(), 1):
        tabs = list(manager)
        active = manager.active_tab
        if not tabs or active is None:
            continue
        count = len(tabs)
        title = get_os_window_title(os_window_id) or f"Window {index}"
        label = f"- {title}: {count} {'tab' if count == 1 else 'tabs'}"
        if current in tabs:
            label += " (current)"
        tab_ids = [tab.id for tab in tabs]
        entries.append(dict(kind='window', label=label, tab_id=active.id,
                            tabs=tab_ids, title=title))
        for tab_index, tab in enumerate(tabs, 1):
            branch = "└─>" if tab_index == count else "├─>"
            marker = "*" if tab is active else ""
            label = f"{branch} {tab_index}: {tab.name or tab.title}{marker}"
            if tab.num_window_groups > 1:
                label += f" [{tab.num_window_groups} panes]"
            entries.append(dict(kind='tab', label=label, tab_id=tab.id,
                                title=tab.name or tab.title))
            snapshots[str(tab.id)] = snapshot_tab(tab, boss)
    selected = None

    def done(answer, target_window_id, boss):
        nonlocal selected
        selected = answer

    # Kittens overlay one pane. Temporarily use stack to fill the tab, keeping
    # the original layout object (including its split tree and size biases).
    layout_state = None
    chooser_layout = None
    if current is not None and current.current_layout.name != 'stack':
        layout_state = (current.current_layout, current._current_layout_name,
                        current._last_used_layout)

    def restore_layout():
        if (layout_state is not None and current in boss.all_tabs
                and current.current_layout is chooser_layout):
            current.current_layout, current._current_layout_name, current._last_used_layout = layout_state
            current.mark_tab_bar_dirty()
            current.relayout()

    def removed(target_window_id, boss):
        restore_layout()
        chosen(selected)

    try:
        if layout_state is not None:
            current._set_current_layout('stack')
            chooser_layout = current.current_layout
            current.relayout()
        overlay = boss.run_kitten_with_metadata(
            "window_picker.py", args=tuple(args[1:]),
            input_data=json.dumps(dict(entries=entries, snapshots=snapshots)),
            window=source,
            custom_callback=done, action_on_removal=removed,
        )
        if overlay is None:
            restore_layout()
    except Exception:
        restore_layout()
        raise