molecule-core/workspace-server/internal/handlers/templates.go

package handlers

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/Molecule-AI/molecule-monorepo/platform/internal/db"
	"github.com/Molecule-AI/molecule-monorepo/platform/internal/provisioner"
	"github.com/docker/docker/client"
	"github.com/gin-gonic/gin"
	"gopkg.in/yaml.v3"
)

// allowedRoots are the container paths that the Files API can browse.
var allowedRoots = map[string]bool{
	"/configs":   true,
	"/workspace": true,
	"/home":      true,
	"/plugins":   true,
}

// maxUploadFiles limits the number of files in a single import/replace.
const maxUploadFiles = 200

type TemplatesHandler struct {
	configsDir string
	docker     *client.Client
}

func NewTemplatesHandler(configsDir string, dockerCli *client.Client) *TemplatesHandler {
	return &TemplatesHandler{configsDir: configsDir, docker: dockerCli}
}

type templateSummary struct {
	ID          string   `json:"id"`
	Name        string   `json:"name"`
	Description string   `json:"description"`
	Tier        int      `json:"tier"`
	Model       string   `json:"model"`
	Skills      []string `json:"skills"`
	SkillCount  int      `json:"skill_count"`
}

// resolveTemplateDir finds the template directory for a workspace on the host.
// Only resolves to actual templates (not ws-* dirs since those are now Docker volumes).
// Returns empty string if no matching template is found.
func (h *TemplatesHandler) resolveTemplateDir(wsName string) string {
	nameDir := filepath.Join(h.configsDir, normalizeName(wsName))
	if _, err := os.Stat(nameDir); err == nil {
		return nameDir
	}
	// Search templates by config.yaml name field (e.g., org-pm has name: "PM")
	if tmpl := findTemplateByName(h.configsDir, wsName); tmpl != "" {
		return filepath.Join(h.configsDir, tmpl)
	}
	return ""
}

// validateRelPath checks that a relative path doesn't escape the target directory.
func validateRelPath(relPath string) error {
	clean := filepath.Clean(relPath)
	if filepath.IsAbs(clean) || strings.HasPrefix(clean, "..") {
		return fmt.Errorf("path traversal blocked: %s", relPath)
	}
	return nil
}

// List handles GET /templates
func (h *TemplatesHandler) List(c *gin.Context) {
	entries, err := os.ReadDir(h.configsDir)
	if err != nil {
		c.JSON(http.StatusOK, []templateSummary{})
		return
	}

	templates := make([]templateSummary, 0)
	for _, entry := range entries {
		if !entry.IsDir() {
			continue
		}

		configPath := filepath.Join(h.configsDir, entry.Name(), "config.yaml")
		data, err := os.ReadFile(configPath)
		if err != nil {
			continue
		}

		var raw struct {
			Name        string   `yaml:"name"`
			Description string   `yaml:"description"`
			Tier        int      `yaml:"tier"`
			Model       string   `yaml:"model"`
			Skills      []string `yaml:"skills"`
		}
		if err := yaml.Unmarshal(data, &raw); err != nil {
			continue
		}

		templates = append(templates, templateSummary{
			ID:          entry.Name(),
			Name:        raw.Name,
			Description: raw.Description,
			Tier:        raw.Tier,
			Model:       raw.Model,
			Skills:      raw.Skills,
			SkillCount:  len(raw.Skills),
		})
	}

	c.JSON(http.StatusOK, templates)
}

// ListFiles handles GET /workspaces/:id/files
// Lists files inside the running container's /configs directory (or /workspace, etc.).
// Falls back to host-side config templates directory when container isn't running.
func (h *TemplatesHandler) ListFiles(c *gin.Context) {
	workspaceID := c.Param("id")
	ctx := c.Request.Context()

	// Query params:
	//   ?root=  — base path in container (default: /configs)
	//   ?path=  — subdirectory to list (relative to root, default: "")
	//   ?depth= — max depth to recurse (default: 1, max: 5)
	rootPath := c.DefaultQuery("root", "/configs")
	if !allowedRoots[rootPath] {
		c.JSON(http.StatusBadRequest, gin.H{"error": "root must be one of: /configs, /workspace, /home, /plugins"})
		return
	}
	subPath := c.DefaultQuery("path", "")
	if subPath != "" {
		if err := validateRelPath(subPath); err != nil {
			c.JSON(http.StatusBadRequest, gin.H{"error": "invalid path"})
			return
		}
	}
	depth := 1
	if d := c.Query("depth"); d != "" {
		n, err := strconv.Atoi(d)
		if err != nil || n < 1 || n > 5 {
			c.JSON(http.StatusBadRequest, gin.H{"error": "depth must be 1-5"})
			return
		}
		depth = n
	}
	listPath := rootPath
	if subPath != "" {
		listPath = rootPath + "/" + subPath
	}

	var wsName string
	if err := db.DB.QueryRowContext(ctx, `SELECT name FROM workspaces WHERE id = $1`, workspaceID).Scan(&wsName); err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "workspace not found"})
		return
	}

	type fileEntry struct {
		Path string `json:"path"`
		Size int64  `json:"size"`
		Dir  bool   `json:"dir"`
	}

	// Try container filesystem first
	if containerName := h.findContainer(ctx, workspaceID); containerName != "" {
		// Portable file listing: works on both GNU and BusyBox/Alpine.
		// Uses find + sh -c stat to output TYPE|SIZE|PATH per line.
		output, err := h.execInContainer(ctx, containerName, []string{
			"sh", "-c",
			fmt.Sprintf(`find '%s' -maxdepth %d -not -path '*/.git/*' -not -path '*/__pycache__/*' -not -path '*/node_modules/*' -not -name .DS_Store | while IFS= read -r f; do
				rel="${f#'%s'/}"; [ "$rel" = '%s' ] && continue; [ -z "$rel" ] && continue
				if [ -d "$f" ]; then echo "d|0|$rel"; else s=$(stat -c %%s "$f" 2>/dev/null || stat -f %%z "$f" 2>/dev/null || echo 0); echo "f|$s|$rel"; fi
			done`, listPath, depth, listPath, listPath),
		})
		if err != nil {
			log.Printf("Container file list failed, falling back to host: %v", err)
		} else {
			var files []fileEntry
			for _, line := range strings.Split(output, "\n") {
				parts := strings.SplitN(line, "|", 3)
				if len(parts) != 3 || parts[2] == "" {
					continue
				}
				size, _ := strconv.ParseInt(parts[1], 10, 64)
				files = append(files, fileEntry{
					Path: parts[2],
					Size: size,
					Dir:  parts[0] == "d",
				})
			}
			if files == nil {
				files = []fileEntry{}
			}
			c.JSON(http.StatusOK, files)
			return
		}
	}

	// Fallback: host-side template dir (only for templates, not ws-* workspace volumes)
	configDir := h.resolveTemplateDir(wsName)
	if configDir == "" {
		c.JSON(http.StatusOK, []fileEntry{})
		return
	}

	walkRoot := configDir
	if subPath != "" {
		walkRoot = filepath.Join(configDir, subPath)
	}
	if _, err := os.Stat(walkRoot); os.IsNotExist(err) {
		c.JSON(http.StatusOK, []fileEntry{})
		return
	}

	var files []fileEntry
	filepath.Walk(walkRoot, func(path string, info os.FileInfo, err error) error {
		if err != nil || path == walkRoot {
			return nil
		}
		rel, _ := filepath.Rel(walkRoot, path)
		// Enforce depth limit
		if strings.Count(rel, string(filepath.Separator))+1 > depth {
			if info.IsDir() {
				return filepath.SkipDir
			}
			return nil
		}
		base := filepath.Base(rel)
		if base == ".git" || base == ".DS_Store" || base == "__pycache__" || base == "node_modules" {
			if info.IsDir() {
				return filepath.SkipDir
			}
			return nil
		}
		files = append(files, fileEntry{
			Path: rel,
			Size: info.Size(),
			Dir:  info.IsDir(),
		})
		return nil
	})

	if files == nil {
		files = []fileEntry{}
	}
	c.JSON(http.StatusOK, files)
}

// ReadFile handles GET /workspaces/:id/files/*path
func (h *TemplatesHandler) ReadFile(c *gin.Context) {
	workspaceID := c.Param("id")
	filePath := c.Param("path")
	if strings.HasPrefix(filePath, "/") {
		filePath = filePath[1:]
	}

	if err := validateRelPath(filePath); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid path"})
		return
	}

	ctx := c.Request.Context()
	rootPath := c.DefaultQuery("root", "/configs")
	if !allowedRoots[rootPath] {
		c.JSON(http.StatusBadRequest, gin.H{"error": "root must be one of: /configs, /workspace, /home, /plugins"})
		return
	}

	var wsName string
	if err := db.DB.QueryRowContext(ctx, `SELECT name FROM workspaces WHERE id = $1`, workspaceID).Scan(&wsName); err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "workspace not found"})
		return
	}

	// Try container first
	if containerName := h.findContainer(ctx, workspaceID); containerName != "" {
		containerPath := rootPath + "/" + filePath
		content, err := h.execInContainer(ctx, containerName, []string{"cat", containerPath})
		if err == nil {
			c.JSON(http.StatusOK, gin.H{
				"path":    filePath,
				"content": content,
				"size":    len(content),
			})
			return
		}
	}

	// Fallback: host-side template dir
	templateDir := h.resolveTemplateDir(wsName)
	if templateDir == "" {
		c.JSON(http.StatusNotFound, gin.H{"error": "file not found (container offline, no template)"})
		return
	}
	// validateRelPath is already called above (line 260) for the container path,
	// but the fallback below uses filePath directly in filepath.Join without
	// any sanitization. Re-validate before the host-side read to close the gap.
	if err := validateRelPath(filePath); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid path"})
		return
	}
	fullPath := filepath.Join(templateDir, filePath)
	data, err := os.ReadFile(fullPath)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "file not found"})
		return
	}

	c.JSON(http.StatusOK, gin.H{
		"path":    filePath,
		"content": string(data),
		"size":    len(data),
	})
}

// WriteFile handles PUT /workspaces/:id/files/*path
func (h *TemplatesHandler) WriteFile(c *gin.Context) {
	workspaceID := c.Param("id")
	filePath := c.Param("path")
	if strings.HasPrefix(filePath, "/") {
		filePath = filePath[1:]
	}

	if err := validateRelPath(filePath); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid path"})
		return
	}

	var body struct {
		Content string `json:"content"`
	}
	if err := c.ShouldBindJSON(&body); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
		return
	}

	ctx := c.Request.Context()
	var wsName string
	if err := db.DB.QueryRowContext(ctx, `SELECT name FROM workspaces WHERE id = $1`, workspaceID).Scan(&wsName); err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "workspace not found"})
		return
	}

	// Write via Docker CopyToContainer when container is running
	if containerName := h.findContainer(ctx, workspaceID); containerName != "" {
		singleFile := map[string]string{filePath: body.Content}
		if err := h.copyFilesToContainer(ctx, containerName, "/configs", singleFile); err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to write file: %v", err)})
			return
		}
		c.JSON(http.StatusOK, gin.H{"status": "saved", "path": filePath})
		return
	}

	// Container offline — write via ephemeral container mounting the config volume
	volName := provisioner.ConfigVolumeName(workspaceID)
	singleFile := map[string]string{filePath: body.Content}
	if err := h.writeViaEphemeral(ctx, volName, singleFile); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to write file: %v", err)})
		return
	}
	c.JSON(http.StatusOK, gin.H{"status": "saved", "path": filePath})
}

// DeleteFile handles DELETE /workspaces/:id/files/*path
func (h *TemplatesHandler) DeleteFile(c *gin.Context) {
	workspaceID := c.Param("id")
	filePath := c.Param("path")
	if strings.HasPrefix(filePath, "/") {
		filePath = filePath[1:]
	}

	if err := validateRelPath(filePath); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid path"})
		return
	}

	ctx := c.Request.Context()
	var wsName string
	if err := db.DB.QueryRowContext(ctx, `SELECT name FROM workspaces WHERE id = $1`, workspaceID).Scan(&wsName); err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "workspace not found"})
		return
	}

	// Delete via docker exec when container is running
	if containerName := h.findContainer(ctx, workspaceID); containerName != "" {
		containerPath := "/configs/" + filePath
		_, err := h.execInContainer(ctx, containerName, []string{"rm", "-rf", containerPath})
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to delete: %v", err)})
			return
		}
		c.JSON(http.StatusOK, gin.H{"status": "deleted", "path": filePath})
		return
	}

	// Container offline — delete via ephemeral container
	volName := provisioner.ConfigVolumeName(workspaceID)
	if err := h.deleteViaEphemeral(ctx, volName, filePath); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to delete: %v", err)})
		return
	}
	c.JSON(http.StatusOK, gin.H{"status": "deleted", "path": filePath})
}

// SharedContext handles GET /workspaces/:id/shared-context
// Returns the files listed in the workspace's config.yaml shared_context field.
func (h *TemplatesHandler) SharedContext(c *gin.Context) {
	workspaceID := c.Param("id")
	ctx := c.Request.Context()

	var wsName string
	if err := db.DB.QueryRowContext(ctx, `SELECT name FROM workspaces WHERE id = $1`, workspaceID).Scan(&wsName); err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "workspace not found"})
		return
	}

	type contextFile struct {
		Path    string `json:"path"`
		Content string `json:"content"`
	}

	// Try reading from running container first
	if containerName := h.findContainer(ctx, workspaceID); containerName != "" {
		configData, err := h.execInContainer(ctx, containerName, []string{"cat", "/configs/config.yaml"})
		if err != nil {
			c.JSON(http.StatusOK, []interface{}{})
			return
		}

		var cfg struct {
			SharedContext []string `yaml:"shared_context"`
		}
		if err := yaml.Unmarshal([]byte(configData), &cfg); err != nil || len(cfg.SharedContext) == 0 {
			c.JSON(http.StatusOK, []interface{}{})
			return
		}

		files := make([]contextFile, 0, len(cfg.SharedContext))
		for _, relPath := range cfg.SharedContext {
			if err := validateRelPath(relPath); err != nil {
				continue
			}
			content, err := h.execInContainer(ctx, containerName, []string{"cat", "/configs/" + relPath})
			if err != nil {
				continue
			}
			files = append(files, contextFile{Path: relPath, Content: content})
		}
		c.JSON(http.StatusOK, files)
		return
	}

	// Fallback to host-side template dir
	configDir := h.resolveTemplateDir(wsName)
	if configDir == "" {
		c.JSON(http.StatusOK, []interface{}{})
		return
	}

	configData, err := os.ReadFile(filepath.Join(configDir, "config.yaml"))
	if err != nil {
		c.JSON(http.StatusOK, []interface{}{})
		return
	}

	var cfg struct {
		SharedContext []string `yaml:"shared_context"`
	}
	if err := yaml.Unmarshal(configData, &cfg); err != nil || len(cfg.SharedContext) == 0 {
		c.JSON(http.StatusOK, []interface{}{})
		return
	}

	files := make([]contextFile, 0, len(cfg.SharedContext))
	for _, relPath := range cfg.SharedContext {
		if err := validateRelPath(relPath); err != nil {
			continue
		}
		data, err := os.ReadFile(filepath.Join(configDir, relPath))
		if err != nil {
			continue
		}
		files = append(files, contextFile{Path: relPath, Content: string(data)})
	}

	c.JSON(http.StatusOK, files)
}