bzzz/sandbox/sandbox.go

package sandbox

import (
	"archive/tar"
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"github.com/anthonyrawlins/bzzz/pkg/config"
	"github.com/anthonyrawlins/bzzz/workspace"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/client"
	"github.com/docker/docker/pkg/stdcopy"
)

// Sandbox represents a stateful, isolated execution environment for a single task.
type Sandbox struct {
	ID            string // The ID of the running container.
	HostPath      string // The path on the host machine mounted as the workspace.
	Workspace     string // The path inside the container that is the workspace.
	dockerCli     *client.Client
	ctx           context.Context
	hcfsWorkspace *workspace.HCFSWorkspace // HCFS-backed workspace
	workspaceManager *workspace.HCFSWorkspaceManager // HCFS workspace manager
}

// CommandResult holds the output of a command executed in the sandbox.
type CommandResult struct {
	StdOut   string
	StdErr   string
	ExitCode int
}

// CreateSandbox provisions a new Docker container for a task.
func CreateSandbox(ctx context.Context, taskImage string, agentConfig *config.AgentConfig) (*Sandbox, error) {
	return CreateSandboxWithHCFS(ctx, taskImage, agentConfig, "", "")
}

// CreateSandboxWithHCFS provisions a new Docker container with HCFS workspace support
func CreateSandboxWithHCFS(ctx context.Context, taskImage string, agentConfig *config.AgentConfig, agentID, taskID string) (*Sandbox, error) {
	if taskImage == "" {
		taskImage = agentConfig.SandboxImage
	}

	// Create a new Docker client
	cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
	if err != nil {
		return nil, fmt.Errorf("failed to create docker client: %w", err)
	}

	// Initialize HCFS workspace manager
	hcfsAPIURL := os.Getenv("HCFS_API_URL")
	if hcfsAPIURL == "" {
		hcfsAPIURL = "http://localhost:8000" // Default HCFS API URL
	}
	
	hcfsMountPath := os.Getenv("HCFS_MOUNT_PATH")
	if hcfsMountPath == "" {
		hcfsMountPath = "/tmp/hcfs-workspaces" // Default mount path
	}
	
	workspaceManager := workspace.NewHCFSWorkspaceManager(hcfsAPIURL, hcfsMountPath)

	var hostPath string
	var hcfsWorkspace *workspace.HCFSWorkspace

	// Create workspace - use HCFS if agent/task IDs provided, fallback to temp dir
	if agentID != "" && taskID != "" {
		// Create HCFS-backed workspace
		hcfsWorkspace, err = workspaceManager.CreateWorkspace(ctx, agentID, taskID, agentConfig)
		if err != nil {
			fmt.Printf("⚠️ Failed to create HCFS workspace, falling back to temp dir: %v\n", err)
			// Fallback to temporary directory
			hostPath, err = os.MkdirTemp("", "bzzz-sandbox-")
			if err != nil {
				return nil, fmt.Errorf("failed to create temp dir for sandbox: %w", err)
			}
		} else {
			hostPath = workspaceManager.GetWorkspacePath(hcfsWorkspace)
		}
	} else {
		// Create a temporary directory on the host (legacy mode)
		hostPath, err = os.MkdirTemp("", "bzzz-sandbox-")
		if err != nil {
			return nil, fmt.Errorf("failed to create temp dir for sandbox: %w", err)
		}
		fmt.Printf("⚠️ Creating sandbox without HCFS (no agent/task ID provided)\n")
	}

	// Read GitHub token for authentication
	githubToken := os.Getenv("BZZZ_GITHUB_TOKEN")
	if githubToken == "" {
		// Try to read from file
		tokenBytes, err := os.ReadFile("/home/tony/chorus/business/secrets/gh-token")
		if err == nil {
			githubToken = strings.TrimSpace(string(tokenBytes))
		}
	}

	// Define container configuration
	containerConfig := &container.Config{
		Image:        taskImage,
		Tty:          true, // Keep the container running
		OpenStdin:    true,
		WorkingDir:   "/home/agent/work",
		User:         "agent",
		Env: []string{
			"GITHUB_TOKEN=" + githubToken,
			"GH_TOKEN=" + githubToken,
		},
	}

	// Define host configuration (e.g., volume mounts, resource limits)
	hostConfig := &container.HostConfig{
		Binds: []string{fmt.Sprintf("%s:/home/agent/work", hostPath)},
		Resources: container.Resources{
			NanoCPUs: 2 * 1000000000, // 2 CPUs
			Memory:   2 * 1024 * 1024 * 1024, // 2GB
		},
	}

	// Create the container
	resp, err := cli.ContainerCreate(ctx, containerConfig, hostConfig, nil, nil, "")
	if err != nil {
		os.RemoveAll(hostPath) // Clean up the directory if container creation fails
		return nil, fmt.Errorf("failed to create container: %w", err)
	}

	// Start the container
	if err := cli.ContainerStart(ctx, resp.ID, container.StartOptions{}); err != nil {
		os.RemoveAll(hostPath) // Clean up
		return nil, fmt.Errorf("failed to start container: %w", err)
	}

	fmt.Printf("✅ Sandbox container %s created successfully.\n", resp.ID[:12])

	return &Sandbox{
		ID:               resp.ID,
		HostPath:         hostPath,
		Workspace:        "/home/agent/work",
		dockerCli:        cli,
		ctx:              ctx,
		hcfsWorkspace:    hcfsWorkspace,
		workspaceManager: workspaceManager,
	}, nil
}

// DestroySandbox stops and removes the container and its associated host directory.
func (s *Sandbox) DestroySandbox() error {
	if s == nil || s.ID == "" {
		return nil
	}

	// Define a timeout for stopping the container
	timeout := 30 // seconds
	
	// Stop the container
	fmt.Printf("🛑 Stopping sandbox container %s...\n", s.ID[:12])
	err := s.dockerCli.ContainerStop(s.ctx, s.ID, container.StopOptions{Timeout: &timeout})
	if err != nil {
		// Log the error but continue to try and clean up
		fmt.Printf("⚠️  Error stopping container %s: %v. Proceeding with cleanup.\n", s.ID, err)
	}

	// Remove the container
	err = s.dockerCli.ContainerRemove(s.ctx, s.ID, container.RemoveOptions{Force: true})
	if err != nil {
		fmt.Printf("⚠️  Error removing container %s: %v. Proceeding with cleanup.\n", s.ID, err)
	}

	// Handle workspace cleanup - HCFS vs temp directory
	if s.hcfsWorkspace != nil && s.workspaceManager != nil {
		// Store any final artifacts before cleanup
		artifacts := s.collectWorkspaceArtifacts()
		if len(artifacts) > 0 {
			if err := s.workspaceManager.StoreWorkspaceArtifacts(s.ctx, s.hcfsWorkspace, artifacts); err != nil {
				fmt.Printf("⚠️ Failed to store workspace artifacts: %v\n", err)
			}
		}
		
		// Destroy HCFS workspace
		if err := s.workspaceManager.DestroyWorkspace(s.ctx, s.hcfsWorkspace); err != nil {
			fmt.Printf("⚠️ Failed to destroy HCFS workspace: %v\n", err)
		}
	} else {
		// Legacy mode: Remove the host directory
		fmt.Printf("🗑️  Removing host directory %s...\n", s.HostPath)
		err = os.RemoveAll(s.HostPath)
		if err != nil {
			return fmt.Errorf("failed to remove host directory %s: %w", s.HostPath, err)
		}
	}

	fmt.Printf("✅ Sandbox %s destroyed successfully.\n", s.ID[:12])
	return nil
}

// RunCommand executes a shell command inside the sandbox.
func (s *Sandbox) RunCommand(command string) (*CommandResult, error) {
	// Update workspace usage if using HCFS
	if s.hcfsWorkspace != nil && s.workspaceManager != nil {
		s.workspaceManager.UpdateWorkspaceUsage(s.hcfsWorkspace)
	}

	// Configuration for the exec process
	execConfig := container.ExecOptions{
		Cmd:          []string{"/bin/sh", "-c", command},
		AttachStdout: true,
		AttachStderr: true,
		Tty:          false,
	}

	// Create the exec instance
	execID, err := s.dockerCli.ContainerExecCreate(s.ctx, s.ID, execConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to create exec in container: %w", err)
	}

	// Start the exec process
	resp, err := s.dockerCli.ContainerExecAttach(s.ctx, execID.ID, container.ExecStartOptions{})
	if err != nil {
		return nil, fmt.Errorf("failed to attach to exec in container: %w", err)
	}
	defer resp.Close()

	// Read the output
	var stdout, stderr bytes.Buffer
	_, err = stdcopy.StdCopy(&stdout, &stderr, resp.Reader)
	if err != nil {
		return nil, fmt.Errorf("failed to read exec output: %w", err)
	}

	// Inspect the exec process to get the exit code
	inspect, err := s.dockerCli.ContainerExecInspect(s.ctx, execID.ID)
	if err != nil {
		return nil, fmt.Errorf("failed to inspect exec in container: %w", err)
	}

	return &CommandResult{
		StdOut:   stdout.String(),
		StdErr:   stderr.String(),
		ExitCode: inspect.ExitCode,
	}, nil
}

// WriteFile writes content to a file inside the sandbox's workspace.
func (s *Sandbox) WriteFile(path string, content []byte) error {
	// Create a temporary file on the host
	tmpfile, err := os.CreateTemp("", "bzzz-write-")
	if err != nil {
		return fmt.Errorf("failed to create temp file: %w", err)
	}
	defer os.Remove(tmpfile.Name())

	if _, err := tmpfile.Write(content); err != nil {
		return fmt.Errorf("failed to write to temp file: %w", err)
	}
	tmpfile.Close()

	// Copy the file into the container
	dstPath := filepath.Join(s.Workspace, path)
	
	// Create tar archive of the file
	tarBuf := new(bytes.Buffer)
	tw := tar.NewWriter(tarBuf)
	
	fileInfo, err := os.Stat(tmpfile.Name())
	if err != nil {
		return fmt.Errorf("failed to stat temp file: %w", err)
	}
	
	header := &tar.Header{
		Name: filepath.Base(path),
		Size: fileInfo.Size(),
		Mode: 0644,
	}
	
	if err := tw.WriteHeader(header); err != nil {
		return fmt.Errorf("failed to write tar header: %w", err)
	}
	
	fileContent, err := os.ReadFile(tmpfile.Name())
	if err != nil {
		return fmt.Errorf("failed to read temp file: %w", err)
	}
	
	if _, err := tw.Write(fileContent); err != nil {
		return fmt.Errorf("failed to write to tar: %w", err)
	}
	
	if err := tw.Close(); err != nil {
		return fmt.Errorf("failed to close tar writer: %w", err)
	}
	
	return s.dockerCli.CopyToContainer(s.ctx, s.ID, filepath.Dir(dstPath), tarBuf, container.CopyToContainerOptions{})
}

// ReadFile reads the content of a file from the sandbox's workspace.
func (s *Sandbox) ReadFile(path string) ([]byte, error) {
	srcPath := filepath.Join(s.Workspace, path)

	// Copy the file from the container
	reader, _, err := s.dockerCli.CopyFromContainer(s.ctx, s.ID, srcPath)
	if err != nil {
		return nil, fmt.Errorf("failed to copy from container: %w", err)
	}
	defer reader.Close()

	// The result is a tar archive, so we need to extract it
	tr := tar.NewReader(reader)
	if _, err := tr.Next(); err != nil {
		return nil, fmt.Errorf("failed to get tar header: %w", err)
	}

	buf := new(bytes.Buffer)
	if _, err := io.Copy(buf, tr); err != nil {
		return nil, fmt.Errorf("failed to read file content from tar: %w", err)
	}

	return buf.Bytes(), nil
}

// collectWorkspaceArtifacts collects important artifacts from the workspace
func (s *Sandbox) collectWorkspaceArtifacts() map[string]string {
	artifacts := make(map[string]string)
	
	// Common artifact files to collect
	artifactFiles := []string{
		"output/result.txt",
		"output/summary.md",
		"logs/execution.log",
		"build/manifest.json",
		".bzzz-workspace-state",
	}
	
	for _, artifactFile := range artifactFiles {
		content, err := s.ReadFile(artifactFile)
		if err == nil && len(content) > 0 {
			artifacts[artifactFile] = string(content)
		}
	}
	
	return artifacts
}

// GetHCFSWorkspace returns the HCFS workspace if available
func (s *Sandbox) GetHCFSWorkspace() *workspace.HCFSWorkspace {
	return s.hcfsWorkspace
}

// IsUsingHCFS returns true if this sandbox is using HCFS workspace
func (s *Sandbox) IsUsingHCFS() bool {
	return s.hcfsWorkspace != nil && s.workspaceManager != nil
}