import os
import shlex
import subprocess

from .logging import logger
from .osworld_types import OSWorldProcessResult, OSWorldProcessSuccessResult
from .util import assert_linux


def execute(command: str | list[str], shell: bool | None) -> OSWorldProcessResult:
    logger.info(f"executing {command} shell: {shell}")

    shell = shell or False
    if isinstance(command, str):
        if shell:
            command = [command]
        else:
            logger.warning(
                "Command should be a list of strings. Now it is a string. Will split it by space."
            )
            command = shlex.split(command)

    # Expand user directory
    for i, arg in enumerate(command):
        if arg.startswith("~/"):
            command[i] = os.path.expanduser(arg)

    # Execute the command without any safety checks.
    assert_linux()
    run_result = subprocess.run(
        command,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        shell=shell,
        text=True,
        timeout=120,
        check=False,
    )

    result = OSWorldProcessSuccessResult(
        status="success",
        output=run_result.stdout,
        error=run_result.stderr,
        returncode=run_result.returncode,
    )

    logger.debug(f"execute of {command} succeeded w/{result}")
    return result
