"""Interactive MCP shell (REPL).

Persistent MCP session, tab completion, command history, @name prefixing,
use <name> live switching, and always-live discovery inside the REPL.
"""

import asyncio
import json
import logging
import sys
from contextlib import AsyncExitStack

import click
from prompt_toolkit import PromptSession
from prompt_toolkit.completion import Completer, Completion
from prompt_toolkit.history import FileHistory

from mcpbundles.args import parse_tool_args, ArgParseError
from mcpbundles.auth_resolve import resolve_auth, resolve_auth_for_named
from mcpbundles.config import AuthError, CLISettings, CONFIG_DIR
from mcpbundles.connections import ConnectionManager
from mcpbundles.mcp_client import hub_session, call_tool, ToolCallError, ConnectionError
from mcpbundles.output import (
    render_tools_table, render_call_result, render_error,
    render_success, render_resources_table, render_prompts_table,
    stdout_console, console,
)
from mcpbundles.schema import render_tool_detail

logger = logging.getLogger(__name__)

HISTORY_FILE = CONFIG_DIR / "shell_history"

SHELL_COMMANDS = {
    "tools": "List available tools",
    "resources": "List available resources",
    "prompts": "List available prompts",
    "call": "Call a tool: call <name> [key=value ...]",
    "read": "Read a resource: read <uri>",
    "exec": "Execute code: exec <code>",
    "use": "Switch connection: use <name>",
    "connections": "List active connections",
    "help": "Show available commands",
    "exit": "Exit the shell",
    "quit": "Exit the shell",
}


class ShellCompleter(Completer):
    """Dynamic completer that fetches tool names from the session."""

    def __init__(self) -> None:
        self._tool_names: list[str] = []

    def update_tools(self, names: list[str]) -> None:
        self._tool_names = names

    def get_completions(self, document, complete_event):
        text = document.text_before_cursor
        words = text.split()

        if len(words) == 0 or (len(words) == 1 and not text.endswith(" ")):
            prefix = words[0] if words else ""
            for cmd in SHELL_COMMANDS:
                if cmd.startswith(prefix):
                    yield Completion(cmd, start_position=-len(prefix), display_meta=SHELL_COMMANDS[cmd])
        elif len(words) >= 1 and words[0] in ("call", "tools") and (len(words) == 1 and text.endswith(" ") or len(words) == 2 and not text.endswith(" ")):
            prefix = words[1] if len(words) > 1 else ""
            for name in self._tool_names:
                if name.startswith(prefix):
                    yield Completion(name, start_position=-len(prefix))
        elif len(words) >= 1 and words[0] == "use" and (len(words) == 1 and text.endswith(" ") or len(words) == 2 and not text.endswith(" ")):
            prefix = words[1] if len(words) > 1 else ""
            mgr = ConnectionManager()
            for conn in mgr.list_connections():
                if conn.name.startswith(prefix):
                    yield Completion(conn.name, start_position=-len(prefix))


class _ShellSession:
    """Manages the active MCP session inside the REPL, supporting live switching."""

    def __init__(self, completer: ShellCompleter) -> None:
        self._completer = completer
        self._session = None
        self._stack: AsyncExitStack | None = None
        self.conn_name: str = ""
        self.url: str = ""
        self.endpoint: str = ""

    async def connect(self, url: str, endpoint: str, conn_name: str, auth: str) -> None:
        await self.disconnect()
        self.url = url
        self.endpoint = endpoint
        self.conn_name = conn_name
        self._stack = AsyncExitStack()
        try:
            self._session = await self._stack.enter_async_context(
                hub_session(auth, url, endpoint)
            )
            try:
                tools_result = await self._session.list_tools()
                self._completer.update_tools([t.name for t in tools_result.tools])
            except Exception:
                pass
        except Exception:
            await self._cleanup_stack()
            raise

    async def disconnect(self) -> None:
        await self._cleanup_stack()
        self._session = None
        self._completer.update_tools([])

    async def _cleanup_stack(self) -> None:
        if self._stack:
            try:
                await self._stack.aclose()
            except Exception:
                pass
            self._stack = None

    @property
    def session(self):
        if self._session is None:
            raise ConnectionError("No active session. Use 'use <name>' to connect.")
        return self._session


@click.command("shell")
@click.option("--as", "connection", help="Start shell with a named connection")
@click.pass_context
def shell_cmd(ctx: click.Context, connection: str | None) -> None:
    """Start an interactive MCP shell.

    Features:
        - Persistent MCP session (no reconnect per command)
        - Tab completion for commands and tool names
        - Command history across sessions
        - use <name> to switch connections live
    """
    asyncio.run(_shell_impl(ctx, connection))


async def _shell_impl(ctx: click.Context, connection: str | None) -> None:
    settings: CLISettings = ctx.obj["settings"]

    try:
        auth = resolve_auth(settings, connection)
    except (AuthError, SystemExit):
        raise SystemExit(1)

    conn_name = auth.conn_name or "hub"

    HISTORY_FILE.parent.mkdir(parents=True, exist_ok=True)
    completer = ShellCompleter()
    prompt_session: PromptSession = PromptSession(
        history=FileHistory(str(HISTORY_FILE)),
        completer=completer,
    )

    shell = _ShellSession(completer)

    stdout_console.print(f"[bold]MCPBundles Shell[/bold] — connecting to [cyan]{conn_name}[/]...")

    try:
        await shell.connect(auth.url, auth.endpoint, conn_name, auth.token)
    except (AuthError, ConnectionError) as exc:
        render_error(str(exc))
        raise SystemExit(1)

    stdout_console.print(f"Connected. Type [cyan]help[/] for commands, [cyan]exit[/] to quit.\n")

    try:
        while True:
            try:
                line = await asyncio.to_thread(
                    prompt_session.prompt,
                    f"mcpbundles({shell.conn_name})> ",
                )
            except (EOFError, KeyboardInterrupt):
                stdout_console.print("\nGoodbye!")
                return

            line = line.strip()
            if not line:
                continue

            parts = line.split(None, 1)
            cmd = parts[0].lower()
            rest = parts[1] if len(parts) > 1 else ""

            if cmd in ("exit", "quit"):
                stdout_console.print("Goodbye!")
                return
            elif cmd == "help":
                for c, desc in SHELL_COMMANDS.items():
                    stdout_console.print(f"  [cyan]{c:<14}[/] {desc}")
            elif cmd == "tools":
                try:
                    result = await shell.session.list_tools()
                    completer.update_tools([t.name for t in result.tools])
                    if rest:
                        match = next((t for t in result.tools if t.name == rest), None)
                        if match:
                            render_tool_detail(match)
                        else:
                            render_error(f"Tool '{rest}' not found.")
                    else:
                        render_tools_table(result.tools)
                except Exception as exc:
                    render_error(str(exc))
            elif cmd == "resources":
                try:
                    result = await shell.session.list_resources()
                    render_resources_table(result.resources)
                except Exception as exc:
                    render_error(str(exc))
            elif cmd == "prompts":
                try:
                    result = await shell.session.list_prompts()
                    render_prompts_table(result.prompts)
                except Exception as exc:
                    render_error(str(exc))
            elif cmd == "call":
                if not rest:
                    render_error("Usage: call <tool_name> [key=value ...]")
                    continue
                call_parts = rest.split()
                tool_name = call_parts[0]
                arg_parts = tuple(call_parts[1:])
                try:
                    arguments = parse_tool_args(arg_parts)
                except ArgParseError as exc:
                    render_error(str(exc))
                    continue
                try:
                    result = await call_tool(shell.session, tool_name, arguments)
                    render_call_result(result, mode="pretty")
                except ToolCallError as exc:
                    render_error(str(exc))
            elif cmd == "exec":
                if not rest:
                    render_error("Usage: exec <code>")
                    continue
                try:
                    result = await call_tool(shell.session, "code_execution", {"code": rest})
                    render_call_result(result, mode="pretty")
                except ToolCallError as exc:
                    render_error(str(exc))
            elif cmd == "read":
                if not rest:
                    render_error("Usage: read <uri>")
                    continue
                try:
                    result = await shell.session.read_resource(rest)
                    for content in result.contents:
                        if hasattr(content, "text") and content.text:
                            print(content.text)
                except Exception as exc:
                    render_error(str(exc))
            elif cmd == "use":
                if not rest:
                    render_error("Usage: use <connection_name>")
                    continue
                target_name = rest.strip().lstrip("@")
                mgr = ConnectionManager()
                target_conn = mgr.get(target_name)
                if not target_conn:
                    render_error(f"Connection '{target_name}' not found. Use 'connections' to list available.")
                    continue
                try:
                    target_auth = resolve_auth_for_named(target_conn, mgr, settings)
                except AuthError as exc:
                    render_error(str(exc))
                    continue
                console.print(f"  Switching to [cyan]{target_name}[/]...")
                try:
                    await shell.connect(target_conn.url, target_conn.endpoint, target_name, target_auth)
                    render_success(f"Connected to {target_name}")
                except (AuthError, ConnectionError) as exc:
                    render_error(f"Failed to connect to {target_name}: {exc}")
            elif cmd == "connections":
                mgr = ConnectionManager()
                conns = mgr.list_connections()
                if conns:
                    for c in conns:
                        marker = " [green]★[/]" if c.name == shell.conn_name else ""
                        stdout_console.print(f"  [cyan]{c.name}[/] → {c.display_target}{marker}")
                else:
                    stdout_console.print(f"  [cyan]{shell.conn_name}[/] → [dim]{shell.url}{shell.endpoint}[/] [green]★[/]")
            else:
                render_error(f"Unknown command: {cmd}. Type 'help' for available commands.")

    except Exception as exc:
        render_error(f"Session error: {exc}")
        raise SystemExit(1)
    finally:
        await shell.disconnect()
