fetch_ml/internal/jupyter/workspace_manager.go

package jupyter

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/jfraeys/fetch_ml/internal/logging"
)

const (
	mountTypeBind = "bind"
)

// WorkspaceManager handles workspace mounting and volume management for Jupyter services
type WorkspaceManager struct {
	logger   *logging.Logger
	basePath string
	mounts   map[string]*WorkspaceMount
}

// WorkspaceMount represents a workspace mount configuration
type WorkspaceMount struct {
	ID            string            `json:"id"`
	HostPath      string            `json:"host_path"`
	ContainerPath string            `json:"container_path"`
	MountType     string            `json:"mount_type"` // "bind", "volume", "tmpfs"
	ReadOnly      bool              `json:"read_only"`
	Options       map[string]string `json:"options"`
	Services      []string          `json:"services"` // Service IDs using this mount
}

// MountRequest defines parameters for creating a workspace mount
type MountRequest struct {
	HostPath      string            `json:"host_path"`
	ContainerPath string            `json:"container_path"`
	MountType     string            `json:"mount_type"`
	ReadOnly      bool              `json:"read_only"`
	Options       map[string]string `json:"options"`
}

// NewWorkspaceManager creates a new workspace manager
func NewWorkspaceManager(logger *logging.Logger, basePath string) *WorkspaceManager {
	return &WorkspaceManager{
		logger:   logger,
		basePath: basePath,
		mounts:   make(map[string]*WorkspaceMount),
	}
}

// CreateWorkspaceMount creates a new workspace mount
func (wm *WorkspaceManager) CreateWorkspaceMount(req *MountRequest) (*WorkspaceMount, error) {
	// Validate request
	if err := wm.validateMountRequest(req); err != nil {
		return nil, err
	}

	// Generate mount ID
	mountID := wm.generateMountID(req.HostPath)

	// Ensure host path exists
	if req.MountType == mountTypeBind {
		if err := wm.ensureHostPath(req.HostPath); err != nil {
			return nil, fmt.Errorf("failed to prepare host path: %w", err)
		}
	}

	// Create mount record
	mount := &WorkspaceMount{
		ID:            mountID,
		HostPath:      req.HostPath,
		ContainerPath: req.ContainerPath,
		MountType:     req.MountType,
		ReadOnly:      req.ReadOnly,
		Options:       req.Options,
		Services:      []string{},
	}

	// Store mount
	wm.mounts[mountID] = mount

	wm.logger.Info("workspace mount created",
		"mount_id", mountID,
		"host_path", req.HostPath,
		"container_path", req.ContainerPath,
		"mount_type", req.MountType)

	return mount, nil
}

// GetMount retrieves a mount by ID
func (wm *WorkspaceManager) GetMount(mountID string) (*WorkspaceMount, error) {
	mount, exists := wm.mounts[mountID]
	if !exists {
		return nil, fmt.Errorf("mount %s not found", mountID)
	}
	return mount, nil
}

// FindMountByPath finds a mount by host path
func (wm *WorkspaceManager) FindMountByPath(hostPath string) (*WorkspaceMount, error) {
	for _, mount := range wm.mounts {
		if mount.HostPath == hostPath {
			return mount, nil
		}
	}
	return nil, fmt.Errorf("no mount found for path %s", hostPath)
}

// ListMounts returns all mounts
func (wm *WorkspaceManager) ListMounts() []*WorkspaceMount {
	mounts := make([]*WorkspaceMount, 0, len(wm.mounts))
	for _, mount := range wm.mounts {
		mounts = append(mounts, mount)
	}
	return mounts
}

// RemoveMount removes a workspace mount
func (wm *WorkspaceManager) RemoveMount(mountID string) error {
	mount, exists := wm.mounts[mountID]
	if !exists {
		return fmt.Errorf("mount %s not found", mountID)
	}

	// Check if mount is in use
	if len(mount.Services) > 0 {
		return fmt.Errorf("cannot remove mount %s: in use by services %v", mountID, mount.Services)
	}

	// Remove mount
	delete(wm.mounts, mountID)

	wm.logger.Info("workspace mount removed", "mount_id", mountID, "host_path", mount.HostPath)

	return nil
}

// AttachService attaches a service to a mount
func (wm *WorkspaceManager) AttachService(mountID, serviceID string) error {
	mount, exists := wm.mounts[mountID]
	if !exists {
		return fmt.Errorf("mount %s not found", mountID)
	}

	// Check if service is already attached
	for _, service := range mount.Services {
		if service == serviceID {
			return nil // Already attached
		}
	}

	// Attach service
	mount.Services = append(mount.Services, serviceID)

	wm.logger.Debug("service attached to mount",
		"mount_id", mountID,
		"service_id", serviceID)

	return nil
}

// DetachService detaches a service from a mount
func (wm *WorkspaceManager) DetachService(mountID, serviceID string) error {
	mount, exists := wm.mounts[mountID]
	if !exists {
		return fmt.Errorf("mount %s not found", mountID)
	}

	// Find and remove service
	for i, service := range mount.Services {
		if service == serviceID {
			mount.Services = append(mount.Services[:i], mount.Services[i+1:]...)
			wm.logger.Debug("service detached from mount",
				"mount_id", mountID,
				"service_id", serviceID)
			return nil
		}
	}

	return nil // Service not attached
}

// GetMountsForService returns all mounts used by a service
func (wm *WorkspaceManager) GetMountsForService(serviceID string) []*WorkspaceMount {
	var serviceMounts []*WorkspaceMount
	for _, mount := range wm.mounts {
		for _, service := range mount.Services {
			if service == serviceID {
				serviceMounts = append(serviceMounts, mount)
				break
			}
		}
	}
	return serviceMounts
}

// PrepareWorkspace prepares a workspace for Jupyter
func (wm *WorkspaceManager) PrepareWorkspace(workspacePath string) (*WorkspaceMount, error) {
	// Check if mount already exists
	mount, err := wm.FindMountByPath(workspacePath)
	if err == nil {
		return mount, nil
	}

	// Create new mount
	req := &MountRequest{
		HostPath:      workspacePath,
		ContainerPath: "/workspace",
		MountType:     mountTypeBind,
		ReadOnly:      false,
		Options: map[string]string{
			"Z": "", // For SELinux compatibility
		},
	}

	return wm.CreateWorkspaceMount(req)
}

// validateMountRequest validates a mount request
func (wm *WorkspaceManager) validateMountRequest(req *MountRequest) error {
	if req.HostPath == "" {
		return fmt.Errorf("host path is required")
	}

	if req.ContainerPath == "" {
		return fmt.Errorf("container path is required")
	}

	if req.MountType == "" {
		req.MountType = mountTypeBind
	}

	// Validate mount type
	validTypes := []string{"bind", "volume", "tmpfs"}
	valid := false
	for _, t := range validTypes {
		if req.MountType == t {
			valid = true
			break
		}
	}
	if !valid {
		return fmt.Errorf("invalid mount type %s, must be one of: %v", req.MountType, validTypes)
	}

	// For bind mounts, host path must be absolute
	if req.MountType == mountTypeBind && !filepath.IsAbs(req.HostPath) {
		return fmt.Errorf("bind mount host path must be absolute: %s", req.HostPath)
	}

	// Check for duplicate mounts
	for _, mount := range wm.mounts {
		if mount.HostPath == req.HostPath {
			return fmt.Errorf("mount for path %s already exists", req.HostPath)
		}
	}

	return nil
}

// generateMountID generates a unique mount ID
func (wm *WorkspaceManager) generateMountID(hostPath string) string {
	// Create a safe ID from the host path
	safePath := strings.ToLower(hostPath)
	safePath = strings.ReplaceAll(safePath, "/", "-")
	safePath = strings.ReplaceAll(safePath, " ", "-")
	safePath = strings.ReplaceAll(safePath, "_", "-")

	// Remove leading dash
	safePath = strings.TrimPrefix(safePath, "-")

	return fmt.Sprintf("mount-%s", safePath)
}

// ensureHostPath ensures the host path exists and has proper permissions
func (wm *WorkspaceManager) ensureHostPath(hostPath string) error {
	// Check if path exists
	info, err := os.Stat(hostPath)
	if err != nil {
		if os.IsNotExist(err) {
			// Create directory
			if err := os.MkdirAll(hostPath, 0750); err != nil {
				return fmt.Errorf("failed to create directory %s: %w", hostPath, err)
			}
			wm.logger.Info("created workspace directory", "path", hostPath)
		} else {
			return fmt.Errorf("failed to stat path %s: %w", hostPath, err)
		}
	} else if !info.IsDir() {
		return fmt.Errorf("host path %s is not a directory", hostPath)
	}

	// Check permissions
	if err := os.Chmod(hostPath, 0600); err != nil {
		wm.logger.Warn("failed to set permissions on workspace", "path", hostPath, "error", err)
	}

	return nil
}

// ValidateWorkspace validates a workspace for Jupyter use
func (wm *WorkspaceManager) ValidateWorkspace(workspacePath string) error {
	// Check if workspace exists
	info, err := os.Stat(workspacePath)
	if err != nil {
		if os.IsNotExist(err) {
			return fmt.Errorf("workspace %s does not exist", workspacePath)
		}
		return fmt.Errorf("failed to access workspace %s: %w", workspacePath, err)
	}

	if !info.IsDir() {
		return fmt.Errorf("workspace %s is not a directory", workspacePath)
	}

	// Check for common Jupyter files
	jupyterFiles := []string{
		"*.ipynb",
		"requirements.txt",
		"environment.yml",
		"Pipfile",
		"pyproject.toml",
	}

	var foundFiles []string
	for _, pattern := range jupyterFiles {
		matches, err := filepath.Glob(filepath.Join(workspacePath, pattern))
		if err == nil && len(matches) > 0 {
			foundFiles = append(foundFiles, pattern)
		}
	}

	if len(foundFiles) == 0 {
		wm.logger.Warn("workspace may not be a Jupyter project",
			"path", workspacePath,
			"no_files_found", strings.Join(jupyterFiles, ", "))
	} else {
		wm.logger.Info("workspace validated",
			"path", workspacePath,
			"found_files", strings.Join(foundFiles, ", "))
	}

	return nil
}

// GetWorkspaceInfo returns information about a workspace
func (wm *WorkspaceManager) GetWorkspaceInfo(workspacePath string) (*WorkspaceInfo, error) {
	info := &WorkspaceInfo{
		Path: workspacePath,
	}

	// Get directory info
	dirInfo, err := os.Stat(workspacePath)
	if err != nil {
		return nil, fmt.Errorf("failed to stat workspace: %w", err)
	}
	info.Size = dirInfo.Size()
	info.Modified = dirInfo.ModTime()

	// Count files
	err = filepath.Walk(workspacePath, func(path string, fi os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !fi.IsDir() {
			info.FileCount++
			info.TotalSize += fi.Size()

			// Categorize files
			ext := strings.ToLower(filepath.Ext(path))
			switch ext {
			case ".py":
				info.PythonFiles++
			case ".ipynb":
				info.NotebookFiles++
			case ".txt", ".md":
				info.TextFiles++
			case ".json", ".yaml", ".yml":
				info.ConfigFiles++
			}
		}
		return nil
	})

	if err != nil {
		return nil, fmt.Errorf("failed to scan workspace: %w", err)
	}

	return info, nil
}

// WorkspaceInfo contains information about a workspace
type WorkspaceInfo struct {
	Path          string    `json:"path"`
	FileCount     int64     `json:"file_count"`
	PythonFiles   int64     `json:"python_files"`
	NotebookFiles int64     `json:"notebook_files"`
	TextFiles     int64     `json:"text_files"`
	ConfigFiles   int64     `json:"config_files"`
	Size          int64     `json:"size"`
	TotalSize     int64     `json:"total_size"`
	Modified      time.Time `json:"modified"`
}

// Cleanup removes unused mounts
func (wm *WorkspaceManager) Cleanup() error {
	var toRemove []string
	for mountID, mount := range wm.mounts {
		if len(mount.Services) == 0 {
			toRemove = append(toRemove, mountID)
		}
	}

	for _, mountID := range toRemove {
		if err := wm.RemoveMount(mountID); err != nil {
			wm.logger.Warn("failed to remove unused mount", "mount_id", mountID, "error", err)
		}
	}

	if len(toRemove) > 0 {
		wm.logger.Info("cleanup completed", "removed_mounts", len(toRemove))
	}

	return nil
}