from pathlib import Path
import shlex
from typing import Callable, Literal, Self

from textual import on
from textual.reactive import var, Initialize
from textual.app import ComposeResult

from textual.actions import SkipAction
from textual.binding import Binding

from textual.content import Content
from textual import getters
from textual.message import Message
from textual.widgets import OptionList, TextArea, Label
from textual import containers
from textual.widget import Widget
from textual.widgets.option_list import Option
from textual.widgets.text_area import Selection
from textual import events

from toad.app import ToadApp
from toad import messages
from toad.widgets.highlighted_textarea import HighlightedTextArea
from toad.widgets.condensed_path import CondensedPath
from toad.widgets.path_search import PathSearch
from toad.widgets.plan import Plan
from toad.widgets.question import Ask, Question
from toad.widgets.slash_complete import SlashComplete
from toad.messages import UserInputSubmitted
from toad.slash_command import SlashCommand
from toad.prompt.extract import extract_paths_from_prompt
from toad.acp.agent import Mode
from toad.path_complete import PathComplete


class ModeSwitcher(OptionList):
    BINDINGS = [Binding("escape", "dismiss")]

    @on(OptionList.OptionSelected)
    def on_option_selected(self, event: OptionList.OptionSelected):
        self.post_message(messages.ChangeMode(event.option_id))
        self.blur()

    def action_dismiss(self):
        self.blur()


class InvokeFileSearch(Message):
    pass


class InvokeSlashComplete(Message):
    pass


class AgentInfo(Label):
    pass


class ModeInfo(Label):
    pass


class StatusLine(Label):
    status: var[str] = var("")

    def watch_status(self, status: str) -> None:
        self.set_class(not bool(status), "-hidden")
        self.update(status)
        self.tooltip = status


class PromptContainer(containers.HorizontalGroup):
    def on_mouse_down(self, event: events.MouseUp) -> None:
        for child in self.query("*"):
            if child.has_focus:
                return
        prompt_text_area = self.query_one(PromptTextArea)
        if not prompt_text_area.has_focus:
            prompt_text_area.focus()


class PromptTextArea(HighlightedTextArea):
    HELP = """\
## Prompt

This is your prompt, where you can talk to your AI.
See on-screen instructions for details.

- Be simple
- Be direct
- Nothing fancy
"""

    BINDING_GROUP_TITLE = "Prompt"

    BINDINGS = [
        Binding(
            "enter",
            "submit",
            "Send",
            key_display="⏎",
            priority=True,
            tooltip="Send the prompt to the agent",
        ),
        Binding(
            "ctrl+j,shift+enter",
            "newline",
            "Line",
            key_display="⇧+⏎",
            tooltip="Insert a new line character",
        ),
        Binding(
            "ctrl+j,shift+enter",
            "multiline_submit",
            "Send",
            key_display="⇧+⏎",
            tooltip="Send the prompt to the agent",
        ),
        Binding(
            "tab",
            "tab_complete",
            "Complete",
            tooltip="Complete path (if possible)",
            priority=True,
            show=False,
        ),
    ]

    app = getters.app(ToadApp)

    auto_completes: var[list[Option]] = var(list)
    multi_line = var(False, bindings=True)
    shell_mode = var(False, bindings=True)
    agent_ready: var[bool] = var(False)
    path_complete: var[PathComplete] = var(Initialize(lambda obj: PathComplete()))
    suggestions: var[list[str] | None] = var(None)
    suggestions_index: var[int] = var(0)

    project_path = var(Path())
    working_directory = var("")

    slash_commands: var[list[SlashCommand]] = var([])
    slash_command_prefixes: var[tuple[str, ...]] = var(())

    class Submitted(Message):
        def __init__(self, markdown: str) -> None:
            self.markdown = markdown
            super().__init__()

    class RequestShellMode(Message):
        pass

    class CancelShell(Message):
        pass

    def watch_slash_commands(self, slash_commands: list[SlashCommand]) -> None:
        """A tuple of slash commands for performance reasons (used with `str.startswith`)."""
        self.slash_command_prefixes = tuple(
            [slash_command.command for slash_command in slash_commands]
        )

    def highlight_slash_command(self, text: str) -> Content:
        """Override slash command highlighting."""

        if text.startswith(self.slash_command_prefixes):
            content = Content(text)
            for slash_command in self.slash_commands:
                if text.startswith(slash_command.command + " "):
                    content = content.stylize(
                        "$text-success", 0, len(slash_command.command)
                    )
                    break
            return content
        return Content(text)

    def highlight_shell(self, text: str) -> Content:
        """Override shell highlighting with additional danger detection."""
        content = super().highlight_shell(text)

        if not self.app.settings.get("shell.warn_dangerous", bool):
            return content

        from toad import danger

        spans, _danger_level = danger.detect(
            str(self.project_path), self.working_directory, content.plain
        )
        content = content.add_spans(spans)
        return content

    def on_mount(self) -> None:
        self.highlight_cursor_line = False
        self.hide_suggestion_on_blur = False

    def on_key(self, event: events.Key) -> None:
        if (
            not self.shell_mode
            and self.cursor_location == (0, 0)
            and event.character in {"!", "$"}
        ):
            self.post_message(self.RequestShellMode())
            event.prevent_default()
        elif self.shell_mode and event.key == "tab":
            event.prevent_default()
        elif event.key != "escape":
            self.suggestions = None
            self.suggestion = ""

    def update_suggestion(self) -> None:
        prompt = self.query_ancestor(Prompt)

        if self.selection.start == self.selection.end and self.text.startswith("/"):
            return

        if self.shell_mode and self.cursor_at_end_of_text and "\n" not in self.text:
            if prompt.complete_callback is not None:
                if completes := prompt.complete_callback(self.text):
                    if self.text not in completes:
                        self.suggestion = completes[-1]

    def check_action(self, action: str, parameters: tuple[object, ...]) -> bool | None:
        if action == "newline" and self.multi_line:
            return False
        if action == "submit" and self.multi_line:
            return False
        if action == "multiline_submit":
            return self.multi_line
        return True

    def action_multiline_submit(self) -> None:
        if not self.agent_ready:
            self.app.bell()
            self.post_message(
                messages.Flash(
                    "Agent is not ready. Please wait while the agent connects…",
                    "warning",
                )
            )
            return
        self.post_message(UserInputSubmitted(self.text, self.shell_mode))
        self.clear()

    def action_newline(self) -> None:
        self.insert("\n")

    def action_submit(self) -> None:
        if not self.agent_ready and not self.shell_mode:
            self.app.bell()
            self.post_message(
                messages.Flash(
                    "Agent is not ready. Please wait while the agent connects…",
                    "warning",
                )
            )
            return
        if self.suggestion:
            if " " not in self.text:
                self.insert(self.suggestion + " ")
            else:
                prompt = self.query_ancestor(Prompt)
                last_token = shlex.split(self.text + self.suggestion)[-1]
                last_token_path = Path(prompt.working_directory) / last_token
                if last_token_path.is_dir():
                    self.insert(self.suggestion)
                else:
                    self.insert(self.suggestion + " ")
                self.suggestion = ""
            return
        self.post_message(UserInputSubmitted(self.text, self.shell_mode))
        self.clear()

    def action_cursor_up(self, select: bool = False):
        if self.selection.is_empty and not select:
            row, _column = self.selection[0]
            if row == 0:
                self.post_message(messages.HistoryMove(-1, self.shell_mode, self.text))
                return
        super().action_cursor_up(select)

    def action_cursor_down(self, select: bool = False):
        if self.selection.is_empty and not select:
            row, _column = self.selection[0]
            if row == (self.wrapped_document.height - 1):
                self.post_message(messages.HistoryMove(+1, self.shell_mode, self.text))
                return
        super().action_cursor_down(select)

    def action_delete_left(self) -> None:
        selection = self.selection
        if selection.start == selection.end and self.selection.end == (0, 0):
            self.post_message(self.CancelShell())
            return
        return super().action_delete_left()

    async def action_tab_complete(self) -> None:
        if not self.shell_mode:
            return

        import shlex

        prompt = self.query_ancestor(Prompt)

        if not self.cursor_at_end_of_text:
            return

        _cursor_row, cursor_column = prompt.prompt_text_area.selection.end
        pre_complete = self.text[:cursor_column]
        post_complete = self.text[cursor_column:]
        shlex_tokens = shlex.split(pre_complete)
        if not shlex_tokens:
            return

        command = shlex_tokens[0]

        exclude_node_type: Literal["file"] | Literal["dir"] | None = None
        if (
            command
            in self.app.settings.get("shell.directory_commands", str).splitlines()
        ):
            exclude_node_type = "file"
        elif command in self.app.settings.get("shell.file_commands", str).splitlines():
            exclude_node_type = "dir"

        tab_complete, suggestions = await self.path_complete(
            Path(prompt.working_directory),
            shlex_tokens[-1],
            exclude_type=exclude_node_type,
        )

        if tab_complete is not None:
            shlex_tokens = shlex_tokens[:-1] + [shlex_tokens[-1] + tab_complete]
            path_component = Path(prompt.working_directory) / shlex_tokens[-1]
            if path_component.is_file():
                spaces = " "
            else:
                spaces = ""

            self.clear()
            self.insert(
                " ".join(token.replace(" ", "\\ ") for token in shlex_tokens)
                + post_complete
                + spaces
            )
            self.suggestions = None
        else:
            if suggestions != self.suggestions:
                self.suggestions = suggestions or None
                self.suggestions_index = 0
                if suggestions:
                    self.suggestion = suggestions[0]
            elif self.suggestions:
                self.suggestions_index = (self.suggestions_index + 1) % len(
                    self.suggestions
                )
                self.suggestion = self.suggestions[self.suggestions_index]

    async def watch_selection(
        self, previous_selection: Selection, selection: Selection
    ) -> None:
        if previous_selection == selection:
            return
        if selection.start == selection.end:
            previous_y, previous_x = previous_selection.end
            y, x = selection.end
            if y == previous_y:
                direction = -1 if x < previous_x else +1
            else:
                direction = 0
            line = self.document.get_line(y)

            if (
                not self.shell_mode
                and y == 0
                and x == 1
                and direction == +1
                and line
                and line[0] == "/"
            ):
                self.post_message(InvokeSlashComplete())
                return

            if y == 0 and line and line[0] == "/" and direction == -1:
                if line in self.slash_command_prefixes:
                    self.selection = Selection((0, 0), (0, len(line)))
                    return

            for _path, start, end in extract_paths_from_prompt(line):
                if x > start and x < end:
                    self.selection = Selection((y, start), (y, end))
                    break
                if direction == -1 and x == end:
                    self.selection = Selection((y, start), (y, end))
                    break

            if x > 0 and x <= len(line) and line[x - 1] == "@":
                remaining_line = line[x + 1 :]
                if not remaining_line or remaining_line[0].isspace():
                    self.post_message(InvokeFileSearch())


class Prompt(containers.VerticalGroup):

    BINDINGS = [
        Binding("escape", "dismiss", "Dismiss"),
    ]

    PROMPT_NULL = " "
    PROMPT_SHELL = Content.styled("$", "$text-primary")
    PROMPT_AI = Content.styled("❯", "$text-secondary")
    PROMPT_MULTILINE = Content.styled("☰", "$text-secondary")

    prompt_container = getters.query_one("#prompt-container", Widget)
    prompt_text_area = getters.query_one(PromptTextArea)
    prompt_label = getters.query_one("#prompt", Label)
    current_directory = getters.query_one(CondensedPath)
    path_search = getters.query_one(PathSearch)
    slash_complete = getters.query_one(SlashComplete)
    question = getters.query_one(Question)
    mode_switcher = getters.query_one(ModeSwitcher)

    slash_commands: var[list[SlashCommand]] = var(list)
    shell_mode = var(False)
    multi_line = var(False)
    show_path_search = var(False, toggle_class="-show-path-search", bindings=True)
    show_slash_complete = var(False, toggle_class="-show-slash-complete", bindings=True)
    project_path = var(Path())
    working_directory = var("")
    agent_info = var(Content(""))
    _ask: var[Ask | None] = var(None)
    plan: var[list[Plan.Entry]]
    agent_ready: var[bool] = var(False)
    current_mode: var[Mode | None] = var(None)
    modes: var[dict[str, Mode] | None] = var(None)
    status: var[str] = var("")

    app = getters.app(ToadApp)

    def __init__(
        self,
        project_path: Path,
        *,
        name: str | None = None,
        id: str | None = None,
        classes: str | None = None,
        disabled: bool = False,
        complete_callback: Callable[[str], list[str]] | None = None,
    ):
        super().__init__(name=name, id=id, classes=classes, disabled=disabled)
        self.ask_queue: list[Ask] = []
        self.complete_callback = complete_callback
        self.project_path = project_path

    @property
    def text(self) -> str:
        return self.prompt_text_area.text

    @text.setter
    def text(self, text: str) -> None:
        self.prompt_text_area.text = text
        self.prompt_text_area.selection = Selection.cursor(
            self.prompt_text_area.get_cursor_line_end_location()
        )

    def watch_current_mode(self, mode: Mode | None) -> None:
        self.set_class(mode is not None, "-has-mode")
        if mode is not None:
            tooltip = Content.from_markup(
                "[b]$description[/]\n\n[dim](click to open mode switcher)",
                description=mode.description,
            )
            self.query_one(ModeInfo).with_tooltip(tooltip).update(mode.name)
        self.watch_modes(self.modes)

    async def watch_project_path(self) -> None:
        """Initial refresh of paths."""
        self.call_later(self.path_search.refresh_paths)

    def ask(self, ask: Ask) -> None:
        """Replace the textarea prompt with a menu of options.

        Args:
            ask: An `Ask` instance which contains a question and responses.
        """
        self.ask_queue.append(ask)
        self.app.terminal_alert()
        if self._ask is None:
            self._ask = self.ask_queue.pop(0)

    @on(events.Click, "ModeInfo")
    def on_click(self):
        self.mode_switcher.focus()

    def watch_modes(self, modes: dict[str, Mode] | None) -> None:
        from toad.visuals.columns import Columns

        columns = Columns("auto", "auto", "flex")
        if modes is not None:
            mode_list = sorted(modes.values(), key=lambda mode: mode.name.lower())
            for mode in mode_list:
                columns.add_row(
                    (
                        Content.styled("✔", "$text-success")
                        if self.current_mode and mode.id == self.current_mode.id
                        else ""
                    ),
                    Content.from_markup("[bold]$mode[/]", mode=mode.name),
                    Content.styled(mode.description or "", "dim"),
                )
        else:
            mode_list = []

        self.mode_switcher.set_options(
            [Option(row, id=mode.id) for row, mode in zip(columns, mode_list)]
        )
        if self.current_mode is not None:
            self.mode_switcher.highlighted = self.mode_switcher.get_option_index(
                self.current_mode.id
            )

    def watch_agent_ready(self, ready: bool) -> None:
        self.set_class(not ready, "-not-ready")
        if ready:
            self.query_one(AgentInfo).update(self.agent_info)

    def watch_agent_info(self, agent_info: Content) -> None:
        if self.agent_ready:
            self.query_one(AgentInfo).update(agent_info)
        else:
            self.query_one(AgentInfo).update("Initializing…")

    def watch_multiline(self) -> None:
        self.update_prompt()

    def watch_shell_mode(self) -> None:
        self.update_prompt()

    def watch_working_directory(self, working_directory: str) -> None:
        if not working_directory:
            return
        out_of_bounds = not Path(working_directory).is_relative_to(self.project_path)
        if out_of_bounds and not self.has_class("-working-directory-out-of-bounds"):
            self.post_message(
                messages.Flash(
                    "You have navigated away from the project directory",
                    style="error",
                    duration=5,
                )
            )
        self.set_class(
            out_of_bounds,
            "-working-directory-out-of-bounds",
        )

    def watch__ask(self, ask: Ask | None) -> None:
        self.set_class(ask is not None, "-mode-ask")
        if ask is None:
            self.prompt_text_area.focus()
        else:
            self.question.update(ask)
            self.question.focus()

    def update_prompt(self):
        """Update the prompt according to the current mode."""
        if self.shell_mode:
            self.prompt_label.update(self.PROMPT_SHELL, layout=False)
            self.add_class("-shell-mode")
            self.prompt_text_area.placeholder = Content.from_markup(
                "Enter shell command\t[r]▌esc▐[/r] prompt mode"
            ).expand_tabs(8)
            self.prompt_text_area.highlight_language = "shell"
        else:
            self.prompt_label.update(
                self.PROMPT_MULTILINE if self.multi_line else self.PROMPT_AI,
                layout=False,
            )
            self.remove_class("-shell-mode")

            self.prompt_text_area.placeholder = Content.assemble(
                "What would you like to do?\t".expandtabs(8),
                ("▌!▐", "r"),
                " shell ",
                ("▌/▐", "r"),
                " commands ",
                ("▌@▐", "r"),
                " files",
            )
            self.prompt_text_area.highlight_language = "markdown"

    @property
    def likely_shell(self) -> bool:
        text = self.prompt_text_area.text
        if "\n" in text or " " in text or not text.strip():
            return False

        shell_commands = {
            command.strip()
            for command in self.app.settings.get(
                "shell.allow_commands", expect_type=str
            ).split()
        }
        if text.split(" ", 1)[0] in shell_commands:
            return True
        return False

    @property
    def is_shell_mode(self) -> bool:
        return self.shell_mode or self.likely_shell

    def focus(self, scroll_visible: bool = True) -> Self:
        if self._ask is not None:
            self.question.focus()
        else:
            self.query(HighlightedTextArea).focus()
        return self

    def append(self, text: str) -> None:
        self.query_one(HighlightedTextArea).insert(
            text, maintain_selection_offset=False
        )

    def watch_show_path_search(self, show: bool) -> None:
        self.prompt_text_area.suggestion = ""

    def watch_show_slash_complete(self, show: bool) -> None:
        self.slash_complete.focus()

    def project_directory_updated(self) -> None:
        """Called when there is may be new files"""
        self.path_search.refresh_paths()

    @on(PromptTextArea.RequestShellMode)
    def on_request_shell_mode(self, event: PromptTextArea.RequestShellMode):
        self.shell_mode = True
        self.update_prompt()

    @on(TextArea.Changed)
    def on_text_area_changed(self, event: TextArea.Changed) -> None:
        text = event.text_area.text

        self.multi_line = "\n" in text or "```" in text

        if not self.multi_line and self.likely_shell:
            self.shell_mode = True

        self.update_prompt()

    @on(PromptTextArea.CancelShell)
    def on_cancel_shell(self, event: PromptTextArea.CancelShell):
        self.shell_mode = False

    @on(InvokeFileSearch)
    def on_invoke_file_search(self, event: InvokeFileSearch) -> None:
        event.stop()
        if not self.shell_mode:
            self.show_path_search = True
            self.path_search.reset()

    @on(InvokeSlashComplete)
    def on_invoke_slash_complete(self, event: InvokeSlashComplete) -> None:
        event.stop()
        self.show_slash_complete = True

    @on(messages.PromptSuggestion)
    def on_prompt_suggestion(self, event: messages.PromptSuggestion) -> None:
        event.stop()
        self.prompt_text_area.suggestion = event.suggestion

    @on(SlashComplete.Completed)
    def on_slash_complete_completed(self, event: SlashComplete.Completed) -> None:
        self.prompt_text_area.clear()
        self.prompt_text_area.insert(f"{event.command} ")

    @on(messages.Dismiss)
    def on_dismiss(self, event: messages.Dismiss) -> None:
        event.stop()
        if event.widget is self.slash_complete:
            self.show_slash_complete = False
            self.prompt_text_area.suggestion = ""
            self.focus()
        elif event.widget is self.path_search:
            self.show_path_search = False
            self.focus()

    @on(messages.InsertPath)
    def on_insert_path(self, event: messages.InsertPath) -> None:
        event.stop()
        if " " in event.path:
            path = f'"{event.path}"'
        else:
            path = event.path
            if (
                self.prompt_text_area.get_text_range(*self.prompt_text_area.selection)
                != " "
            ):
                path += " "
        self.prompt_text_area.insert(path)

    @on(Question.Answer)
    def on_question_answer(self, event: Question.Answer) -> None:
        """Question has been answered."""
        event.stop()

        def remove_question() -> None:
            """Remove the question and restore the text prompt."""
            if self.ask_queue:
                self._ask = self.ask_queue.pop(0)
            else:
                self._ask = None
            self.app.terminal_alert(False)

        if self._ask is not None and (callback := self._ask.callback) is not None:
            callback(event.answer)

        self.set_timer(0.3, remove_question)

    def suggest(self, suggestion: str) -> None:
        if suggestion.startswith(self.text) and self.text != suggestion:
            self.prompt_text_area.suggestion = suggestion[len(self.text) :]

    def compose(self) -> ComposeResult:
        yield PathSearch(self.project_path).data_bind(root=Prompt.project_path)
        yield SlashComplete().data_bind(slash_commands=Prompt.slash_commands)
        with PromptContainer(id="prompt-container"):
            yield Question()
            with containers.HorizontalGroup(id="text-prompt"):
                yield Label(self.PROMPT_AI, id="prompt", markup=False)
                yield PromptTextArea().data_bind(
                    multi_line=Prompt.multi_line,
                    shell_mode=Prompt.shell_mode,
                    agent_ready=Prompt.agent_ready,
                    project_path=Prompt.project_path,
                    working_directory=Prompt.working_directory,
                    slash_commands=Prompt.slash_commands,
                )
        with containers.HorizontalGroup(id="info-container"):
            yield AgentInfo()
            yield CondensedPath().data_bind(path=Prompt.working_directory)
            yield StatusLine(markup=False).data_bind(status=Prompt.status)
            yield ModeSwitcher()
            yield ModeInfo("mode")

    def check_action(self, action: str, parameters: tuple[object, ...]) -> bool | None:
        return True

    def action_dismiss(self) -> None:
        if self.prompt_text_area.suggestion:
            self.prompt_text_area.suggestion = ""
            return
        if self.shell_mode:
            self.shell_mode = False
        elif self.show_slash_complete:
            self.show_slash_complete = False
        else:
            raise SkipAction()
