import asyncio

from SimpleLLMFunc import tool
from .common import print_tool_output


EXECUTE_COMMAND_TIMEOUT_SECONDS = 600


def _build_command_failure_message(result) -> str:
    parts = [f"Command failed with exit code {result.returncode}."]

    stdout = result.stdout.strip()
    stderr = result.stderr.strip()

    if stdout:
        parts.append(f"STDOUT:\n{stdout}")
    if stderr:
        parts.append(f"STDERR:\n{stderr}")

    return "\n\n".join(parts)


def _build_command_timeout_message(exc) -> str:
    parts = [
        f"Command timed out after {EXECUTE_COMMAND_TIMEOUT_SECONDS} seconds.",
        "The process may be stuck, waiting for input, or simply taking too long.",
    ]

    stdout = (exc.stdout or "").strip()
    stderr = (exc.stderr or "").strip()

    if stdout:
        parts.append(f"Partial STDOUT:\n{stdout}")
    if stderr:
        parts.append(f"Partial STDERR:\n{stderr}")

    return "\n\n".join(parts)


@tool(
    name="execute_command",
    description="Execute a system command in shell and return the output.",
)
async def execute_command(command: str) -> str:
    """Execute a system command in shell and return the output.

    Args:
        command: The system command to execute, recommended commands are uv run python <script path>
    Returns:
        The command output (stdout on success, stderr on failure)
    """
    import subprocess
    import time

    try:
        print_tool_output("⚡ Running Command", f"Executing: {command}")

        start_time = time.time()
        result = await asyncio.to_thread(
            subprocess.run,
            command,
            shell=True,
            capture_output=True,
            text=True,
            timeout=EXECUTE_COMMAND_TIMEOUT_SECONDS,
        )

        execution_time = time.time() - start_time

        if result.returncode == 0:
            print_tool_output(
                "✅ Command Completed",
                f"Return code: {result.returncode}, Time: {execution_time:.2f}s, Output: {len(result.stdout)} chars",
            )
            return result.stdout.strip()
        else:
            print_tool_output(
                "❌ Command Failed",
                f"Command failed.\nError: {result.stderr.strip()}",
            )

            return _build_command_failure_message(result)

    except subprocess.TimeoutExpired as exc:
        print_tool_output(
            "⏱️ Command Timed Out", f"Timeout while executing command: {str(exc)}"
        )
        return _build_command_timeout_message(exc)

    except Exception as e:
        print_tool_output("💥 Command Error", f"Command execution failed: {str(e)}")
        return f"Command execution failed: {str(e)}"