"""BashTool20250124 - Persistent bash session with async buffer polling.

Based on Anthropic's reference implementation from anthropic-quickstarts.
Uses asyncio.create_subprocess_shell + direct buffer reads instead of
threaded readline, which avoids exhausting the default ThreadPoolExecutor.
"""

import asyncio
import os
from typing import Any, Literal

from nanobot.agent.tools.anthropic.base import BaseAnthropicTool, ToolResult, ToolError


class _BashSession:
    """A session of a bash shell.

    Uses asyncio subprocess with direct buffer polling — no threads.
    Based on anthropics/anthropic-quickstarts computer-use-demo.
    """

    command: str = "/bin/bash"
    _output_delay: float = 0.2  # seconds between buffer polls
    _timeout: float = 120.0  # seconds
    _sentinel: str = "<<exit>>"

    def __init__(self):
        self._started = False
        self._timed_out = False
        self._process: asyncio.subprocess.Process | None = None

    async def start(self):
        if self._started:
            return

        self._process = await asyncio.create_subprocess_shell(
            self.command,
            preexec_fn=os.setsid,
            shell=True,
            bufsize=0,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        self._started = True

    def stop(self):
        """Terminate the bash shell."""
        if not self._started:
            return
        if self._process and self._process.returncode is None:
            self._process.terminate()

    async def run(self, command: str) -> ToolResult:
        """Execute a command in the bash shell."""
        if not self._started:
            raise ToolError("Session has not started.")
        if self._process is None or self._process.returncode is not None:
            return ToolResult(
                system="tool must be restarted",
                error=f"bash has exited with returncode "
                      f"{self._process.returncode if self._process else 'unknown'}",
            )
        if self._timed_out:
            raise ToolError(
                f"timed out: bash has not returned in {self._timeout} seconds "
                "and must be restarted",
            )

        assert self._process.stdin
        assert self._process.stdout
        assert self._process.stderr

        # Send command + sentinel on its own line so heredoc terminators
        # aren't corrupted (EOF; echo '...' ≠ EOF)
        self._process.stdin.write(
            command.encode() + f"\necho '{self._sentinel}'\n".encode()
        )
        await self._process.stdin.drain()

        # Poll stdout buffer until sentinel appears — no threads involved
        try:
            async with asyncio.timeout(self._timeout):
                while True:
                    await asyncio.sleep(self._output_delay)
                    output = self._process.stdout._buffer.decode()
                    if self._sentinel in output:
                        output = output[: output.index(self._sentinel)]
                        break
        except asyncio.TimeoutError:
            self._timed_out = True
            raise ToolError(
                f"timed out: bash has not returned in {self._timeout} seconds "
                "and must be restarted",
            ) from None

        if output.endswith("\n"):
            output = output[:-1]

        error = self._process.stderr._buffer.decode()
        if error.endswith("\n"):
            error = error[:-1]

        # Clear buffers for next command
        self._process.stdout._buffer.clear()
        self._process.stderr._buffer.clear()

        # Return as ToolResult (our loop handles this type)
        if error and output:
            return ToolResult(output=f"{output}\n\nstderr: {error}")
        elif error:
            return ToolResult(output=error)
        else:
            return ToolResult(output=output if output else "(no output)")


class BashTool20250124(BaseAnthropicTool):
    """Anthropic's native bash_20250124 tool with persistent session.

    Executes bash commands in a long-running shell session. Environment
    variables and working directory persist across commands.

    Parameters:
        command (str, optional): Bash command to execute
        restart (bool, optional): Restart the bash session (clears state)
    """

    api_type: Literal["bash_20250124"] = "bash_20250124"
    name: Literal["bash"] = "bash"
    beta_flag: str | None = None

    def __init__(self):
        self._session: _BashSession | None = None

    async def __call__(
        self,
        command: str | None = None,
        restart: bool = False,
        **kwargs: Any,
    ) -> ToolResult:
        if restart:
            if self._session:
                self._session.stop()
            self._session = _BashSession()
            await self._session.start()
            return ToolResult(system="tool has been restarted.")

        if self._session is None:
            self._session = _BashSession()
            await self._session.start()

        if command is not None:
            try:
                return await self._session.run(command)
            except ToolError as e:
                return ToolResult(error=str(e))

        return ToolResult(error="Either 'command' or 'restart=True' must be provided.")

    def to_params(self) -> dict[str, Any]:
        return {
            "type": self.api_type,
            "name": self.name,
        }
