import os
import shlex
import uuid
from pathlib import Path

from ._shell import AsyncShell, ShellStatus
from pantheon.toolset import ToolSet, tool
from pantheon.utils.log import logger
from pantheon.internal.package_runtime.context import build_context_env


class ShellToolSet(ToolSet):
    """Shell toolset for running shell commands.

    Args:
        name: The name of the toolset.
        **kwargs: Additional keyword arguments.
    """

    def __init__(
        self,
        name: str,
        workdir: str | None = None,
        **kwargs,
    ):
        super().__init__(name, **kwargs)
        self.clientid_to_shellid = {}
        self.shells = {}
        self.workdir = Path(workdir).expanduser().resolve() if workdir else Path.cwd()

    def _prepare_shell_env(self) -> dict:
        return os.environ.copy()

    def _current_context_dict(self) -> dict:
        ctx = dict(self.get_context() or {})
        return ctx

    def _build_runtime_env(self) -> dict:
        effective_workdir = self._get_effective_workdir() or str(self.workdir)
        return build_context_env(
            workdir=effective_workdir,
            context_variables=self._current_context_dict(),
        )

    async def _apply_context_to_shell(self, shell: AsyncShell):
        # cd to effective workdir if set
        effective_workdir = self._get_effective_workdir()
        if effective_workdir:
            try:
                await shell.run_command(f"cd {shlex.quote(effective_workdir)}")
            except Exception:
                logger.warning("Failed to cd to workdir in shell session")

        env = self._build_runtime_env()
        if not env:
            return
        exports = " && ".join(
            f"export {key}={shlex.quote(value)}" for key, value in env.items()
        )
        try:
            await shell.run_command(exports)
        except Exception:
            logger.warning("Failed to propagate context variables to shell session")

    @tool(exclude=True)
    async def new_shell(self) -> dict:
        """Create a new shell and return its id.
        Use `run_command` to run commands."""
        shell = AsyncShell()
        shell.env = self._prepare_shell_env()
        initial_output = await shell.start()
        shell_id = str(uuid.uuid4())
        self.shells[shell_id] = shell

        # Show shell creation status
        logger.info(f"[dim]New shell created: {shell_id[:8]}[/dim]")

        return {
            "success": True,
            "shell_id": shell_id,
            "initial_output": initial_output,
        }

    @tool
    async def close_shell(self, shell_id: str) -> dict:
        """Close a shell.

        Args:
            shell_id: The id of the shell to close.
        """
        shell = self.shells.get(shell_id)
        if shell is None:
            return {"success": False, "error": "Shell not found", "shell_id": shell_id}

        await shell.close()
        del self.shells[shell_id]

        # Show shell closure status
        logger.info(f"[dim]Shell closed: {shell_id[:8]}[/dim]")
        return {"success": True, "shell_id": shell_id}

    @tool
    async def get_shell_output(
        self,
        shell_id: str,
        timeout: int = 5,
        max_output: int | None = None,
    ) -> dict:
        """Get output from a shell, used to check status of background commands.

        When a command times out, it continues running in the background.
        Use this tool with the shell_id returned from run_command to fetch
        the remaining output and check if the command has completed.

        Args:
            shell_id: The shell ID returned from run_command when it timed out.
            timeout: Seconds to wait for output. Default 5 seconds.
            max_output: Optional maximum length for the output. If the output exceeds this,
                it will be smartly truncated (head + tail).

        Returns:
            dict: {
                "output": str,  # Command output since last read
                "status": "completed" | "timeout"  # Whether command finished
            }
        """
        result = await self.run_command_in_shell(
            shell_id=shell_id,
            command=None,
            timeout=timeout,
        )

        # Apply max_output truncation
        truncated = False
        if max_output and result.get("success") and result.get("output"):
            output = result["output"]
            if len(output) > max_output:
                from pantheon.utils.truncate import truncate_string
                result["output"] = truncate_string(output, max_output)
                truncated = True
        
        # Add truncated flag (always present)
        if "success" in result and result["success"]:
            result["truncated"] = truncated

        return result

    @tool(exclude=True)
    async def run_command_in_shell(
        self,
        shell_id: str,
        command: str | None = None,
        timeout: int | None = None,
    ) -> dict:
        """Execute a command or fetch pending output from an existing shell.

        Args:
            shell_id: Identifier for the shell session (from `new_shell` or managed automatically by `run_command`).
            command: Command string to run. If omitted, the tool drains buffered output from the shell (e.g., after a timeout).
            timeout: Optional timeout in seconds. On timeout, the command keeps running and you can re-call this tool with
                only the shell_id to fetch remaining output.

        Returns:
            dict: Structured result including success flag, `status` ("completed" or "timeout"), and the raw output text.
        """

        shell = self.shells.get(shell_id)
        if shell is None:
            return {"success": False, "error": "Shell not found", "shell_id": shell_id}

        status = "completed"

        try:

            def _handle_timeout(message: str) -> None:
                nonlocal status, output
                status = "timeout"
                logger.info(f"[yellow]⚠️ {message}[/yellow]")
                output += "\n[Warning] The execution of the command was interrupted because of the timeout. "

            if command is not None:
                await self._apply_context_to_shell(shell)
                output, finished = await shell.run_command(command, timeout=timeout)
                if timeout is not None and not finished:
                    _handle_timeout(
                        f"Command timed out after {timeout}s - call run_command without a command to fetch remaining output"
                    )
            else:
                marker = shell.current_marker if shell.current_marker else None
                output, finished = await shell.read_until_marker(
                    marker, timeout=timeout
                )
                if finished and shell.current_marker:
                    shell.status = ShellStatus.IDLE
                    shell.current_marker = None

                if timeout is not None and not finished:
                    _handle_timeout(f"Reading shell output timed out after {timeout}s")
        except Exception as exc:
            return {"success": False, "error": str(exc), "shell_id": shell_id}

        response = {
            "success": True,
            "shell_id": shell_id,
            "status": status,
            "output": output,
        }
        if command is not None:
            response["command"] = command
        return response

    def _is_shell_alive(self, shell_id: str) -> bool:
        """Check if a shell is still alive and responsive."""
        if shell_id not in self.shells:
            return False

        shell = self.shells[shell_id]
        if not shell.process:
            return False

        # Check if process is still running
        if shell.process.returncode is not None:
            return False

        return True

    async def _get_available_shell(self) -> str:
        """Get an available (idle) shell, creating a new one if necessary."""
        # First, try to find an idle shell
        for shell_id, shell in self.shells.items():
            if shell.is_idle() and self._is_shell_alive(shell_id):
                return shell_id

        # No idle shell available, create a new one
        result = await self.new_shell()
        return result["shell_id"]

    async def _restart_shell(self, client_id: str) -> str:
        """Restart a shell for a given client_id."""
        old_shell_id = self.clientid_to_shellid.get(client_id)

        # Clean up old shell if it exists
        if old_shell_id and old_shell_id in self.shells:
            try:
                await self.close_shell(old_shell_id)
            except Exception:
                pass  # Ignore cleanup errors
            finally:
                if old_shell_id in self.shells:
                    del self.shells[old_shell_id]

        # Create new shell
        logger.warning(f"Shell crashed (client_id: {client_id}), restarting...")
        res = await self.new_shell()
        new_shell_id = res["shell_id"]
        self.clientid_to_shellid[client_id] = new_shell_id
        logger.info(f"Shell restarted (client_id: {client_id})")

        return new_shell_id

    @tool
    async def run_command(
        self,
        command: str | None = None,
        shell_id: str | None = None,
        timeout: int | None = None,
        max_output: int | None = None,
    ):
        """Run a shell command and return the result.

        This tool automatically manages shell sessions. Just provide the `command`
        to execute. Environment variables and working directory are preserved
        across commands in the same session.

        **Timeout Behavior**: When a command times out, it continues running in
        the background. The response includes a `shell_id` which you can use with
        `get_shell_output` to check progress and retrieve the remaining output.

        **Auto Shell Allocation**: If the current shell is busy (running a
        background command), a new shell is automatically allocated.

        Args:
            command: The command to run.
            timeout: Optional timeout in seconds. If the command doesn't complete
                within this time, it continues running in the background.
            shell_id: Optional. Specify a particular shell ID to use.
                If not provided, automatically uses an available shell.
            max_output: Optional. Max output characters (for verbose commands like
                R package install, npm install). Use smaller values to save tokens.

        Returns:
            dict: {
                "success": bool,       # True if executed successfully
                "output": str,         # Command output (stdout + stderr)
                "status": str,         # "completed" or "timeout"
                "shell_id": str        # Shell ID (returned on timeout for follow-up)
            }
        
        Note:
            Large outputs are auto-truncated. Use max_output for early truncation.

        Tips:
            - You may want to limit the length of output for commands that usually rely on paging and may contain very long output (e.g. `ls -R`, `git log`, use `head -n 20` or `git log -n 5`).
            - Use `max_output` parameter as a safety net.

        Examples:
            # Run a quick command
            run_command(command="ls -la")

            # Run a long task with timeout (becomes background)
            run_command(command="npm run build", timeout=10)
            # Returns: {"status": "timeout", "shell_id": "xxx", ...}

            # Run verbose command with early truncation
            run_command(command="R -e 'BiocManager::install(...)'", max_output=5000)
        """
        # If shell_id is provided, use it directly (Manual Mode)
        if shell_id:
            result = await self.run_command_in_shell(
                shell_id=shell_id,
                command=command,
                timeout=timeout,
            )
        else:
            # Auto Mode (Client ID based)
            context_dict = dict(self.get_context() or {})
            client_id = context_dict.get("client_id")
            if client_id is None:
                client_id = "default"
                # Debug level - not a real problem, just informational
                logger.debug("No client id provided, using default client id.")

            initial_output = ""
            # Resolve shell_id from client_id mapping
            _mapped_shell_id = self.clientid_to_shellid.get(client_id)

            # Check if we need to create a new shell
            if (_mapped_shell_id is None) or (_mapped_shell_id not in self.shells):
                res = await self.new_shell()
                _mapped_shell_id = res["shell_id"]
                initial_output = res["initial_output"]
                self.clientid_to_shellid[client_id] = _mapped_shell_id

            # Check if shell is still alive before running command
            if not self._is_shell_alive(_mapped_shell_id):
                _mapped_shell_id = await self._restart_shell(client_id)
                initial_output = ""  # New shell will have its own initial output

            # If mapped shell is busy, get an available shell
            mapped_shell = self.shells.get(_mapped_shell_id)
            if mapped_shell and not mapped_shell.is_idle():
                _mapped_shell_id = await self._get_available_shell()

            result = await self.run_command_in_shell(
                shell_id=_mapped_shell_id,
                command=command,
                timeout=timeout,
            )

            # If the shell crashed, restart it but do not rerun the command automatically
            if not result.get("success") and self._should_restart(result.get("error")):
                logger.warning(
                    f"Shell command failed for shell {_mapped_shell_id[:8]}: {result.get('error')}"
                )
                _mapped_shell_id = await self._restart_shell(client_id)  # Update mapping
                return result

            if not result.get("success"):
                return result

            # Prepend initial output from new shell creation if applicable
            if initial_output and result.get("success"):
                combined_output = result.get("output") or ""
                result["output"] = (
                    f"{initial_output}\n{combined_output}"
                    if combined_output
                    else initial_output
                )

        # Apply early truncation if max_output specified
        if max_output and result.get("success") and result.get("output"):
            output = result["output"]
            if len(output) > max_output:
                from pantheon.utils.truncate import truncate_string
                result["output"] = truncate_string(output, max_output)
                result["truncated"] = True
            else:
                result["truncated"] = False
        elif result.get("success"):
            # No max_output specified
            result["truncated"] = False

        return result

    def _should_restart(self, error_message: str | None) -> bool:
        if not error_message:
            return False
        error_msg = error_message.lower()
        return any(
            keyword in error_msg
            for keyword in ["broken", "closed", "process", "pipe", "not found"]
        )

    async def run_setup(self):
        """Setup the toolset before running it."""
        logger.warning(
            "This ToolSet is not secure, it can be used to execute arbitrary code."
            " Please be careful when using it."
            " Highly recommend using it in a controlled environment like a docker container."
        )
