"""Interactive REPL for invoking MCP tools during dev server sessions."""

from __future__ import annotations

import asyncio
import json
import readline  # noqa: F401 — imported for side-effect (input history)
import shlex
from typing import TYPE_CHECKING

from rich.json import JSON as RichJSON
from rich.panel import Panel
from rich.table import Table

from hatchdx.devserver import logger as dev_logger
from hatchdx.harness.simulator import SimulatorError
from hatchdx.utils.console import console

if TYPE_CHECKING:
    from hatchdx.devserver.engine import DevServerEngine
    from hatchdx.harness.simulator import ToolInfo


# ---------------------------------------------------------------------------
# Input parsing
# ---------------------------------------------------------------------------


def parse_tool_input(line: str) -> tuple[str, dict]:
    """Parse a REPL line into ``(tool_name, arguments)``.

    Accepted formats::

        tool_name                          → ("tool_name", {})
        tool_name {"key": "value"}         → ("tool_name", {"key": "value"})
        tool_name key=value key2=123       → ("tool_name", {"key": "value", "key2": "123"})

    Raises ``ValueError`` on malformed input.
    """
    line = line.strip()
    if not line:
        raise ValueError("Empty input")

    # Find where the tool name ends
    first_space = line.find(" ")
    if first_space == -1:
        return line, {}

    tool_name = line[:first_space]
    rest = line[first_space:].strip()

    if not rest:
        return tool_name, {}

    # Try JSON first
    if rest.startswith(("{", "[")):
        try:
            args = json.loads(rest)
            if not isinstance(args, dict):
                raise ValueError(f"Arguments must be a JSON object, got {type(args).__name__}")
            return tool_name, args
        except json.JSONDecodeError as exc:
            raise ValueError(f"Invalid JSON arguments: {exc}") from exc

    # Fall back to key=value pairs
    args: dict = {}
    try:
        tokens = shlex.split(rest)
    except ValueError as exc:
        raise ValueError(f"Could not parse arguments: {exc}") from exc

    for token in tokens:
        if "=" not in token:
            raise ValueError(
                f"Expected key=value pair, got: {token!r}\n"
                "Use JSON format for complex values: tool_name {\"key\": \"value\"}"
            )
        key, _, value = token.partition("=")
        # Try to parse as JSON literal (numbers, booleans, null)
        try:
            args[key] = json.loads(value)
        except (json.JSONDecodeError, ValueError):
            args[key] = value

    return tool_name, args


# ---------------------------------------------------------------------------
# Tab completion
# ---------------------------------------------------------------------------


class _ToolCompleter:
    """readline completer that tab-completes tool names."""

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

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

    def complete(self, text: str, state: int) -> str | None:
        if state == 0:
            if text:
                self._matches = [n for n in self._tool_names if n.startswith(text)]
            else:
                self._matches = list(self._tool_names)
        return self._matches[state] if state < len(self._matches) else None


# ---------------------------------------------------------------------------
# REPL
# ---------------------------------------------------------------------------


HELP_TEXT = """\
[bold]Commands:[/]
  [cyan]tool_name {{"key": "value"}}[/]  Call a tool with JSON arguments
  [cyan]tool_name key=value[/]          Call a tool with key=value pairs
  [cyan]tool_name[/]                     Call a tool with no arguments

[bold]Special commands:[/]
  [cyan]tools[/]   List available tools
  [cyan]help[/]    Show this help message
  [cyan]quit[/]    Exit the REPL (or press Ctrl+D)\
"""


class DevRepl:
    """Interactive REPL for calling MCP tools on the running dev server."""

    def __init__(self, engine: DevServerEngine) -> None:
        self._engine = engine
        self._completer = _ToolCompleter()

    async def run(self) -> None:
        """Main REPL loop.  Blocks until the user types ``quit`` or Ctrl+D."""
        self._setup_readline()
        self._print_tools()
        console.print()

        loop = asyncio.get_running_loop()

        while True:
            try:
                line: str = await loop.run_in_executor(None, self._read_input)
            except (EOFError, KeyboardInterrupt):
                console.print()
                break

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

            if line in ("quit", "exit"):
                break
            if line == "help":
                console.print(HELP_TEXT)
                continue
            if line == "tools":
                self._print_tools()
                continue

            await self._execute(line)

    # -- Input ----------------------------------------------------------------

    def _read_input(self) -> str:
        """Blocking input read (runs in a thread via ``run_in_executor``)."""
        return input("mcp> ")

    def _setup_readline(self) -> None:
        """Configure tab completion and history."""
        readline.set_completer(self._completer.complete)
        readline.parse_and_bind("tab: complete")
        readline.set_completer_delims(" ")
        self._completer.update_tools(self._engine.tools)

    # -- Execution ------------------------------------------------------------

    async def _execute(self, line: str) -> None:
        """Parse input, call the tool, and display the result."""
        try:
            tool_name, arguments = parse_tool_input(line)
        except ValueError as exc:
            console.print(f"[red]Parse error:[/] {exc}")
            return

        simulator = self._engine.simulator
        if simulator is None:
            console.print("[yellow]Server is not running. Waiting for restart...[/]")
            return

        # Check that the tool exists
        known_tools = self._engine.tools
        if known_tools and tool_name not in known_tools:
            console.print(f"[yellow]Unknown tool:[/] {tool_name}")
            console.print(f"[dim]Available: {', '.join(known_tools)}[/]")
            return

        dev_logger.log_tool_call(tool_name, arguments)

        import time

        t0 = time.perf_counter()
        try:
            result = await simulator.call_tool(tool_name, arguments)
        except SimulatorError as exc:
            dev_logger.log_tool_result(tool_name, 0, is_error=True)
            console.print(f"[red]Error:[/] {exc}")
            return

        latency_ms = (time.perf_counter() - t0) * 1000
        dev_logger.log_tool_result(tool_name, latency_ms, is_error=result.is_error)

        self._print_result(result.content, is_error=result.is_error)

    # -- Display helpers ------------------------------------------------------

    def _print_tools(self) -> None:
        """Print the list of available tools from the engine."""
        tools = self._engine.tools
        self._completer.update_tools(tools)

        if not tools:
            console.print("[yellow]No tools available.[/]")
            return

        table = Table(
            title="Available Tools",
            show_header=True,
            header_style="bold",
            padding=(0, 1),
        )
        table.add_column("Tool", style="cyan")
        table.add_column("Description", style="dim")

        # Try to get full tool info from the simulator
        tool_infos = self._get_tool_infos()
        if tool_infos:
            for info in tool_infos:
                table.add_row(info.name, info.description or "")
        else:
            for name in tools:
                table.add_row(name, "")

        console.print(table)

    def _get_tool_infos(self) -> list[ToolInfo]:
        """Return cached tool info from the last list_tools call, if available."""
        # The engine only stores names; we can't get full info synchronously.
        # A future enhancement could cache ToolInfo objects on the engine.
        return []

    @staticmethod
    def _print_result(content: list[dict], *, is_error: bool = False) -> None:
        """Pretty-print tool call result content blocks."""
        border_style = "red" if is_error else "green"

        for block in content:
            block_type = block.get("type", "unknown")

            if block_type == "text":
                text = block.get("text", "")
                # Try to render as JSON if it looks like JSON
                if text.strip().startswith(("{", "[")):
                    try:
                        json.loads(text)  # validate
                        console.print(
                            Panel(RichJSON(text), border_style=border_style, expand=False)
                        )
                        continue
                    except (json.JSONDecodeError, ValueError):
                        pass
                console.print(
                    Panel(text, border_style=border_style, expand=False)
                )

            elif block_type == "image":
                mime = block.get("mimeType", "unknown")
                data_len = len(block.get("data", ""))
                console.print(
                    f"[dim]Image ({mime}, {data_len} bytes base64)[/]"
                )

            else:
                # Unknown block type — dump as JSON
                console.print(
                    Panel(RichJSON(json.dumps(block)), border_style=border_style, expand=False)
                )

    def refresh_tools(self) -> None:
        """Update the completer with the latest tool list from the engine."""
        self._completer.update_tools(self._engine.tools)
