tasks/command_engine.py

"""
Universal 24/7 Command Engine & Multi-Project Sandboxing
Production-grade execution of arbitrary bash commands with process tracking,
timeout protection, and multi-project workspace isolation.
"""

import os
import sys
import time
import shutil
import asyncio
import logging
import psutil
from typing import Dict, List, Optional, Any
from pydantic import BaseModel, Field

logger = logging.getLogger("antigravity.command_engine")

# Base directory configuration
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PROJECTS_ROOT = os.environ.get("PROJECTS_DIR", os.path.join(BASE_DIR, "projects"))
os.makedirs(PROJECTS_ROOT, exist_ok=True)


class CommandResult(BaseModel):
    command: str
    project_id: str
    exit_code: int
    stdout: str
    stderr: str
    duration_seconds: float
    timestamp: float
    is_success: bool


class ProjectMeta(BaseModel):
    project_id: str
    name: str
    path: str
    size_bytes: int
    file_count: int
    created_at: float
    updated_at: float


class FileItem(BaseModel):
    name: str
    relative_path: str
    is_directory: bool
    size_bytes: int
    modified_at: float


class RunningProcessTracker:
    """Tracks active subprocesses to enable cancellation and monitoring."""
    def __init__(self):
        self._active_processes: Dict[str, asyncio.subprocess.Process] = {}

    def register(self, task_id: str, process: asyncio.subprocess.Process):
        self._active_processes[task_id] = process

    def unregister(self, task_id: str):
        self._active_processes.pop(task_id, None)

    async def terminate(self, task_id: str, timeout: float = 3.0) -> bool:
        process = self._active_processes.get(task_id)
        if not process:
            return False

        try:
            if process.returncode is None:
                process.terminate()
                try:
                    await asyncio.wait_for(process.wait(), timeout=timeout)
                except asyncio.TimeoutError:
                    process.kill()
                    await process.wait()
            return True
        except Exception as e:
            logger.error(f"Error terminating process {task_id}: {e}")
            return False
        finally:
            self._active_processes.pop(task_id, None)


# Global process tracker instance
process_tracker = RunningProcessTracker()


def get_project_dir(project_id: str) -> str:
    """Sanitize and resolve project sandbox directory."""
    # Sanitize project_id to avoid path traversal
    safe_id = "".join(c for c in project_id if c.isalnum() or c in ("-", "_", ".")).strip(".")
    if not safe_id:
        safe_id = "default"
    project_path = os.path.abspath(os.path.join(PROJECTS_ROOT, safe_id))
    
    # Ensure resolved path is within PROJECTS_ROOT
    if not project_path.startswith(os.path.abspath(PROJECTS_ROOT)):
        project_path = os.path.join(PROJECTS_ROOT, "default")

    os.makedirs(project_path, exist_ok=True)
    return project_path


async def run_bash_command(
    command: str,
    project_id: str = "default",
    task_id: Optional[str] = None,
    timeout_seconds: int = 600,
    custom_env: Optional[Dict[str, str]] = None,
    log_callback: Optional[Any] = None
) -> CommandResult:
    """
    Executes an arbitrary bash command inside the sandboxed project directory.
    Captures stdout/stderr in real-time and provides timeout & cancellation protection.
    """
    start_time = time.time()
    cwd = get_project_dir(project_id)

    env = os.environ.copy()
    env["HOME"] = os.path.expanduser("~")
    env["PROJECT_DIR"] = cwd
    env["PYTHONUNBUFFERED"] = "1"
    if custom_env:
        env.update(custom_env)

    logger.info(f"[{project_id}] Executing command: {command} (timeout={timeout_seconds}s)")
    if log_callback:
        await log_callback(f"[$] {command}\n")

    try:
        process = await asyncio.create_subprocess_shell(
            command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=cwd,
            env=env,
            shell=True
        )

        if task_id:
            process_tracker.register(task_id, process)

        stdout_chunks: List[str] = []
        stderr_chunks: List[str] = []

        async def read_stream(stream, chunks: List[str], prefix: str = ""):
            while True:
                line = await stream.readline()
                if not line:
                    break
                decoded = line.decode("utf-8", errors="replace")
                chunks.append(decoded)
                if log_callback:
                    await log_callback(f"{prefix}{decoded}")

        # Stream stdout and stderr concurrently with timeout
        try:
            await asyncio.wait_for(
                asyncio.gather(
                    read_stream(process.stdout, stdout_chunks),
                    read_stream(process.stderr, stderr_chunks, prefix="[STDERR] "),
                    process.wait()
                ),
                timeout=float(timeout_seconds)
            )
            exit_code = process.returncode if process.returncode is not None else 0
        except asyncio.TimeoutError:
            logger.warning(f"Command timed out after {timeout_seconds}s: {command}")
            process.kill()
            await process.wait()
            exit_code = -1
            stderr_chunks.append(f"\n[ERROR] Command timed out after {timeout_seconds} seconds.\n")

        duration = round(time.time() - start_time, 3)
        full_stdout = "".join(stdout_chunks)
        full_stderr = "".join(stderr_chunks)

        return CommandResult(
            command=command,
            project_id=project_id,
            exit_code=exit_code,
            stdout=full_stdout,
            stderr=full_stderr,
            duration_seconds=duration,
            timestamp=start_time,
            is_success=(exit_code == 0)
        )

    except Exception as e:
        duration = round(time.time() - start_time, 3)
        err_msg = f"Failed to execute command: {str(e)}"
        logger.error(err_msg, exc_info=True)
        return CommandResult(
            command=command,
            project_id=project_id,
            exit_code=1,
            stdout="",
            stderr=err_msg,
            duration_seconds=duration,
            timestamp=start_time,
            is_success=False
        )
    finally:
        if task_id:
            process_tracker.unregister(task_id)


def list_projects() -> List[ProjectMeta]:
    """Lists all sandboxed project workspaces with metadata and disk footprints."""
    projects: List[ProjectMeta] = []
    if not os.path.exists(PROJECTS_ROOT):
        return projects

    for entry in os.scandir(PROJECTS_ROOT):
        if entry.is_dir():
            project_path = entry.path
            total_size = 0
            file_count = 0
            try:
                for root, _, files in os.walk(project_path):
                    for f in files:
                        fp = os.path.join(root, f)
                        if not os.path.islink(fp):
                            try:
                                total_size += os.path.getsize(fp)
                                file_count += 1
                            except OSError:
                                pass
                stat = entry.stat()
                projects.append(ProjectMeta(
                    project_id=entry.name,
                    name=entry.name,
                    path=project_path,
                    size_bytes=total_size,
                    file_count=file_count,
                    created_at=stat.st_ctime,
                    updated_at=stat.st_mtime
                ))
            except Exception as e:
                logger.warning(f"Error reading project {entry.name}: {e}")

    # Sort most recently updated first
    projects.sort(key=lambda p: p.updated_at, reverse=True)
    return projects


def create_project(project_id: str, name: Optional[str] = None) -> ProjectMeta:
    """Creates a new project sandbox directory."""
    path = get_project_dir(project_id)
    stat = os.stat(path)
    return ProjectMeta(
        project_id=os.path.basename(path),
        name=name or os.path.basename(path),
        path=path,
        size_bytes=0,
        file_count=0,
        created_at=stat.st_ctime,
        updated_at=stat.st_mtime
    )


def list_project_files(project_id: str, relative_subpath: str = "") -> List[FileItem]:
    """Lists files and folders inside a specific project workspace."""
    base_proj_dir = get_project_dir(project_id)
    target_dir = os.path.abspath(os.path.join(base_proj_dir, relative_subpath.lstrip("/")))

    if not target_dir.startswith(base_proj_dir):
        target_dir = base_proj_dir

    if not os.path.exists(target_dir):
        return []

    items: List[FileItem] = []
    try:
        for entry in os.scandir(target_dir):
            try:
                stat = entry.stat()
                rel_path = os.path.relpath(entry.path, base_proj_dir)
                items.append(FileItem(
                    name=entry.name,
                    relative_path=rel_path,
                    is_directory=entry.is_dir(),
                    size_bytes=stat.st_size if not entry.is_dir() else 0,
                    modified_at=stat.st_mtime
                ))
            except OSError:
                continue
    except Exception as e:
        logger.error(f"Failed to list files for project {project_id}: {e}")

    items.sort(key=lambda x: (not x.is_directory, x.name.lower()))
    return items


def delete_project(project_id: str) -> bool:
    """Safely removes a project sandbox directory."""
    if project_id.lower() in ("default", "system"):
        return False
    path = get_project_dir(project_id)
    if os.path.exists(path) and path.startswith(os.path.abspath(PROJECTS_ROOT)):
        shutil.rmtree(path, ignore_errors=True)
        return True
    return False


def get_system_telemetry() -> Dict[str, Any]:
    """Retrieves hardware, CPU, RAM, 50GB disk capacity, and OS telemetry."""
    cpu_percent = psutil.cpu_percent(interval=None)
    cpu_count = psutil.cpu_count(logical=True)
    virtual_mem = psutil.virtual_memory()

    # Hugging Face Spaces free tier gives 50GB disk
    disk_usage = shutil.disk_usage(BASE_DIR)
    total_gb = round(disk_usage.total / (1024 ** 3), 2)
    used_gb = round(disk_usage.used / (1024 ** 3), 2)
    free_gb = round(disk_usage.free / (1024 ** 3), 2)
    disk_percent = round((disk_usage.used / disk_usage.total) * 100, 1)

    boot_time = psutil.boot_time()
    uptime_seconds = int(time.time() - boot_time)

    return {
        "cpu_usage_percent": cpu_percent,
        "cpu_cores": cpu_count,
        "memory_total_mb": round(virtual_mem.total / (1024 ** 2), 1),
        "memory_used_mb": round(virtual_mem.used / (1024 ** 2), 1),
        "memory_free_mb": round(virtual_mem.available / (1024 ** 2), 1),
        "memory_percent": virtual_mem.percent,
        "disk_total_gb": total_gb,
        "disk_used_gb": used_gb,
        "disk_free_gb": free_gb,
        "disk_percent": disk_percent,
        "uptime_seconds": uptime_seconds,
        "active_processes_count": len(process_tracker._active_processes)
    }