path: root/internal/api/elaborate.go

package api

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

const elaborateTimeout = 30 * time.Second

func buildElaboratePrompt(workDir string) string {
	workDirLine := `    "project_dir":     string  — leave empty unless you have a specific reason to set it,`
	if workDir != "" {
		workDirLine = fmt.Sprintf(`    "project_dir":     string  — use %q for tasks that operate on this codebase, empty string otherwise,`, workDir)
	}
	return `You are a task configuration assistant for Claudomator, an AI task runner that executes tasks by running Claude or Gemini as a subprocess.

Your ONLY job is to convert any user request into a Claudomator task JSON object. You MUST always output valid JSON. Never ask clarifying questions. Never explain. Never refuse. Make reasonable assumptions and produce the JSON.

Output ONLY a valid JSON object matching this schema (no markdown fences, no prose, no explanation):

{
  "name":        string   — short imperative title (≤60 chars),
  "description": string   — 1-2 sentence summary,
  "agent": {
    "type":            "claude" | "gemini",
    "model":           string  — "sonnet" for claude, "gemini-2.5-flash-lite" for gemini,
    "instructions":    string  — detailed, step-by-step instructions for the agent. Must end with a "## Acceptance Criteria" section listing measurable conditions that define success. For coding tasks, include TDD requirements (write failing tests first, then implement),
` + workDirLine + `
    "max_budget_usd":  number  — conservative estimate (0.25–5.00),
    "allowed_tools":   array   — every tool the task genuinely needs. Include "Write" if creating files, "Edit" if modifying files, "Read" if reading files, "Bash" for shell/git/test commands, "Grep"/"Glob" for searching.
  },
  "timeout":  string  — e.g. "15m",
  "priority": string  — "normal" | "high" | "low",
  "tags":     array   — relevant lowercase tags
}`
}

// elaboratedTask mirrors the task creation schema for elaboration responses.
type elaboratedTask struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	Agent       elaboratedAgent `json:"agent"`
	Timeout     string          `json:"timeout"`
	Priority    string          `json:"priority"`
	Tags        []string        `json:"tags"`
}

type elaboratedAgent struct {
	Type         string   `json:"type"`
	Model        string   `json:"model"`
	Instructions string   `json:"instructions"`
	ProjectDir   string   `json:"project_dir"`
	MaxBudgetUSD float64  `json:"max_budget_usd"`
	AllowedTools []string `json:"allowed_tools"`
}

// sanitizeElaboratedTask enforces tool completeness and dev practice compliance.
// It modifies t in place, inferring missing tools from instruction keywords and
// appending required sections when they are absent.
func sanitizeElaboratedTask(t *elaboratedTask) {
	lower := strings.ToLower(t.Agent.Instructions)

	// Build current tool set.
	toolSet := make(map[string]bool, len(t.Agent.AllowedTools))
	for _, tool := range t.Agent.AllowedTools {
		toolSet[tool] = true
	}

	// Infer missing tools from instruction keywords.
	type rule struct {
		tool     string
		keywords []string
	}
	rules := []rule{
		{"Write", []string{"create file", "write file", "new file", "write to", "save to", "output to", "generate file", "creates a file", "create a new file"}},
		{"Edit", []string{"edit", "modify", "refactor", "replace", "patch"}},
		{"Read", []string{"read", "inspect", "examine", "look at the file"}},
		{"Bash", []string{"run", "execute", "bash", "shell", "command", "build", "compile", "git", "install", "make"}},
		{"Grep", []string{"search for", "grep", "find in", "locate in"}},
		{"Glob", []string{"find file", "list file", "search file"}},
	}
	for _, r := range rules {
		if toolSet[r.tool] {
			continue
		}
		for _, kw := range r.keywords {
			if strings.Contains(lower, kw) {
				toolSet[r.tool] = true
				break
			}
		}
	}
	// Edit without Read is almost always wrong.
	if toolSet["Edit"] && !toolSet["Read"] {
		toolSet["Read"] = true
	}
	// Rebuild the list only when tools were added.
	if len(toolSet) > len(t.Agent.AllowedTools) {
		tools := make([]string, 0, len(toolSet))
		for tool := range toolSet {
			tools = append(tools, tool)
		}
		sort.Strings(tools)
		t.Agent.AllowedTools = tools
	}

	// Append an acceptance criteria section when none is present.
	if !strings.Contains(lower, "acceptance") &&
		!strings.Contains(lower, "done when") &&
		!strings.Contains(lower, "success criteria") {
		t.Agent.Instructions += "\n\n## Acceptance Criteria\nBefore finishing, verify all stated goals are met, tests pass (if applicable), and no unintended side effects were introduced."
	}

	// Append a TDD reminder for coding tasks that do not already mention tests.
	if (toolSet["Edit"] || toolSet["Write"]) && !strings.Contains(lower, "test") {
		t.Agent.Instructions += "\n\n## Dev Practices\nFollow TDD: write a failing test first, then implement the minimum code to make it pass. Commit all changes before finishing."
	}
}

// claudeJSONResult is the top-level object returned by `claude --output-format json`.
type claudeJSONResult struct {
	Result  string `json:"result"`
	IsError bool   `json:"is_error"`
}

// geminiJSONResult is the top-level object returned by `gemini --output-format json`.
type geminiJSONResult struct {
	Response string `json:"response"`
}

// extractJSON returns the first top-level JSON object found in s, stripping
// surrounding prose or markdown code fences the model may have added.
func extractJSON(s string) string {
	start := strings.Index(s, "{")
	end := strings.LastIndex(s, "}")
	if start == -1 || end == -1 || end < start {
		return s
	}
	return s[start : end+1]
}

func (s *Server) claudeBinaryPath() string {
	if s.elaborateCmdPath != "" {
		return s.elaborateCmdPath
	}
	if s.claudeBinPath != "" {
		return s.claudeBinPath
	}
	return "claude"
}

func (s *Server) geminiBinaryPath() string {
	if s.geminiBinPath != "" {
		return s.geminiBinPath
	}
	return "gemini"
}

func readProjectContext(workDir string) string {
	if workDir == "" {
		return ""
	}
	var sb strings.Builder
	for _, filename := range []string{"CLAUDE.md", ".agent/worklog.md"} {
		path := filepath.Join(workDir, filename)
		if data, err := os.ReadFile(path); err == nil {
			if sb.Len() > 0 {
				sb.WriteString("\n\n")
			}
			sb.WriteString(fmt.Sprintf("--- %s ---\n%s", filename, string(data)))
		}
	}
	return sb.String()
}

func (s *Server) appendRawNarrative(workDir, prompt string) {
	if workDir == "" {
		return
	}
	docsDir := filepath.Join(workDir, "docs")
	if err := os.MkdirAll(docsDir, 0755); err != nil {
		s.logger.Error("elaborate: failed to create docs directory", "error", err, "path", docsDir)
		return
	}
	path := filepath.Join(docsDir, "RAW_NARRATIVE.md")
	f, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		s.logger.Error("elaborate: failed to open RAW_NARRATIVE.md", "error", err, "path", path)
		return
	}
	defer f.Close()

	entry := fmt.Sprintf("\n--- %s ---\n%s\n", time.Now().Format(time.RFC3339), prompt)
	if _, err := f.WriteString(entry); err != nil {
		s.logger.Error("elaborate: failed to write to RAW_NARRATIVE.md", "error", err, "path", path)
	}
}

func (s *Server) elaborateWithClaude(ctx context.Context, workDir, fullPrompt string) (*elaboratedTask, error) {
	cmd := exec.CommandContext(ctx, s.claudeBinaryPath(),
		"-p", fullPrompt,
		"--system-prompt", buildElaboratePrompt(workDir),
		"--output-format", "json",
		"--model", "haiku",
	)

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	// Claude returns exit 1 if rate limited but still outputs JSON with is_error: true
	err := cmd.Run()
	
	output := stdout.Bytes()
	if len(output) == 0 {
		if err != nil {
			return nil, fmt.Errorf("claude failed: %w (stderr: %s)", err, stderr.String())
		}
		return nil, fmt.Errorf("claude returned no output")
	}

	var wrapper claudeJSONResult
	if jerr := json.Unmarshal(output, &wrapper); jerr != nil {
		return nil, fmt.Errorf("failed to parse claude JSON wrapper: %w (output: %s)", jerr, string(output))
	}

	if wrapper.IsError {
		return nil, fmt.Errorf("claude error: %s", wrapper.Result)
	}

	var result elaboratedTask
	if jerr := json.Unmarshal([]byte(extractJSON(wrapper.Result)), &result); jerr != nil {
		return nil, fmt.Errorf("failed to parse elaborated task JSON: %w (result: %s)", jerr, wrapper.Result)
	}

	return &result, nil
}

func (s *Server) elaborateWithGemini(ctx context.Context, workDir, fullPrompt string) (*elaboratedTask, error) {
	combinedPrompt := fmt.Sprintf("%s\n\n%s", buildElaboratePrompt(workDir), fullPrompt)
	cmd := exec.CommandContext(ctx, s.geminiBinaryPath(),
		"-p", combinedPrompt,
		"--output-format", "json",
		"--model", "gemini-2.5-flash-lite",
	)

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf("gemini failed: %w (stderr: %s)", err, stderr.String())
	}

	var wrapper geminiJSONResult
	if err := json.Unmarshal(stdout.Bytes(), &wrapper); err != nil {
		return nil, fmt.Errorf("failed to parse gemini JSON wrapper: %w (output: %s)", err, stdout.String())
	}

	var result elaboratedTask
	if err := json.Unmarshal([]byte(extractJSON(wrapper.Response)), &result); err != nil {
		return nil, fmt.Errorf("failed to parse elaborated task JSON: %w (response: %s)", err, wrapper.Response)
	}

	return &result, nil
}

// elaboratedStorySubtask is a leaf unit within a story task.
type elaboratedStorySubtask struct {
	Name         string `json:"name"`
	Instructions string `json:"instructions"`
}

// elaboratedStoryTask is one independently-deployable unit in a story plan.
type elaboratedStoryTask struct {
	Name         string                   `json:"name"`
	Instructions string                   `json:"instructions"`
	Subtasks     []elaboratedStorySubtask `json:"subtasks"`
}

// elaboratedStoryValidation describes how to verify the story was successful.
type elaboratedStoryValidation struct {
	Type            string   `json:"type"`
	Steps           []string `json:"steps"`
	SuccessCriteria string   `json:"success_criteria"`
}

// elaboratedStory is the full implementation plan produced by story elaboration.
type elaboratedStory struct {
	Name       string                    `json:"name"`
	BranchName string                    `json:"branch_name"`
	Tasks      []elaboratedStoryTask     `json:"tasks"`
	Validation elaboratedStoryValidation `json:"validation"`
}

func buildStoryElaboratePrompt() string {
	return `You are a software architect. Given a goal, analyze the codebase at /workspace and produce a structured implementation plan as JSON.

Output ONLY valid JSON matching this schema:
{
  "name": "story name",
  "branch_name": "story/kebab-case-name",
  "tasks": [
    {
      "name": "task name",
      "instructions": "detailed instructions including file paths and what to change",
      "subtasks": [
        { "name": "subtask name", "instructions": "..." }
      ]
    }
  ],
  "validation": {
    "type": "build|test|smoke",
    "steps": ["step1", "step2"],
    "success_criteria": "what success looks like"
  }
}

Rules:
- Tasks must be independently buildable (each can be deployed alone)
- Subtasks within a task are order-dependent and run sequentially
- Instructions must include specific file paths, function names, and exact changes
- Instructions must end with: git add -A && git commit -m "..." && git push origin <branch>
- Validation should match the scope: small change = build check; new feature = smoke test`
}

func (s *Server) elaborateStoryWithClaude(ctx context.Context, workDir, goal string) (*elaboratedStory, error) {
	cmd := exec.CommandContext(ctx, s.claudeBinaryPath(),
		"-p", goal,
		"--system-prompt", buildStoryElaboratePrompt(),
		"--output-format", "json",
		"--model", "haiku",
	)
	if workDir != "" {
		cmd.Dir = workDir
	}

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err := cmd.Run()

	output := stdout.Bytes()
	if len(output) == 0 {
		if err != nil {
			return nil, fmt.Errorf("claude failed: %w (stderr: %s)", err, stderr.String())
		}
		return nil, fmt.Errorf("claude returned no output")
	}

	var wrapper claudeJSONResult
	if jerr := json.Unmarshal(output, &wrapper); jerr != nil {
		return nil, fmt.Errorf("failed to parse claude JSON wrapper: %w (output: %s)", jerr, string(output))
	}
	if wrapper.IsError {
		return nil, fmt.Errorf("claude error: %s", wrapper.Result)
	}

	var result elaboratedStory
	if jerr := json.Unmarshal([]byte(extractJSON(wrapper.Result)), &result); jerr != nil {
		return nil, fmt.Errorf("failed to parse elaborated story JSON: %w (result: %s)", jerr, wrapper.Result)
	}
	return &result, nil
}

func (s *Server) elaborateStoryWithGemini(ctx context.Context, workDir, goal string) (*elaboratedStory, error) {
	combinedPrompt := fmt.Sprintf("%s\n\n%s", buildStoryElaboratePrompt(), goal)
	cmd := exec.CommandContext(ctx, s.geminiBinaryPath(),
		"-p", combinedPrompt,
		"--output-format", "json",
		"--model", "gemini-2.5-flash-lite",
	)
	if workDir != "" {
		cmd.Dir = workDir
	}

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf("gemini failed: %w (stderr: %s)", err, stderr.String())
	}

	var wrapper geminiJSONResult
	if err := json.Unmarshal(stdout.Bytes(), &wrapper); err != nil {
		return nil, fmt.Errorf("failed to parse gemini JSON wrapper: %w (output: %s)", err, stdout.String())
	}

	var result elaboratedStory
	if err := json.Unmarshal([]byte(extractJSON(wrapper.Response)), &result); err != nil {
		return nil, fmt.Errorf("failed to parse elaborated story JSON: %w (response: %s)", err, wrapper.Response)
	}
	return &result, nil
}

func (s *Server) handleElaborateStory(w http.ResponseWriter, r *http.Request) {
	var input struct {
		Goal      string `json:"goal"`
		ProjectID string `json:"project_id"`
	}
	if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid JSON: " + err.Error()})
		return
	}
	if input.Goal == "" {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "goal is required"})
		return
	}
	if input.ProjectID == "" {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "project_id is required"})
		return
	}

	proj, err := s.store.GetProject(input.ProjectID)
	if err != nil {
		writeJSON(w, http.StatusNotFound, map[string]string{"error": "project not found"})
		return
	}

	// Update git refs without modifying the working tree.
	if proj.LocalPath != "" {
		gitCmd := exec.Command("git", "-C", proj.LocalPath, "fetch", "origin")
		if err := gitCmd.Run(); err != nil {
			s.logger.Warn("story elaborate: git fetch failed", "error", err, "path", proj.LocalPath)
		}
	}

	ctx, cancel := context.WithTimeout(r.Context(), elaborateTimeout)
	defer cancel()

	result, err := s.elaborateStoryWithClaude(ctx, proj.LocalPath, input.Goal)
	if err != nil {
		s.logger.Warn("story elaborate: claude failed, falling back to gemini", "error", err)
		result, err = s.elaborateStoryWithGemini(ctx, proj.LocalPath, input.Goal)
		if err != nil {
			s.logger.Error("story elaborate: fallback gemini also failed", "error", err)
			writeJSON(w, http.StatusBadGateway, map[string]string{
				"error": fmt.Sprintf("elaboration failed: %v", err),
			})
			return
		}
	}

	if result.Name == "" {
		writeJSON(w, http.StatusBadGateway, map[string]string{
			"error": "elaboration failed: missing required fields in response",
		})
		return
	}

	writeJSON(w, http.StatusOK, result)
}

func (s *Server) handleElaborateTask(w http.ResponseWriter, r *http.Request) {
	if s.elaborateLimiter != nil && !s.elaborateLimiter.allow(realIP(r)) {
		writeJSON(w, http.StatusTooManyRequests, map[string]string{"error": "rate limit exceeded"})
		return
	}

	var input struct {
		Prompt    string `json:"prompt"`
		ProjectID string `json:"project_id"`
		// project_dir kept for backward compat; project_id takes precedence
		ProjectDir string `json:"project_dir"`
	}
	if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid JSON: " + err.Error()})
		return
	}
	if input.Prompt == "" {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "prompt is required"})
		return
	}

	workDir := s.workDir
	if input.ProjectID != "" {
		if proj, err := s.store.GetProject(input.ProjectID); err == nil {
			workDir = proj.LocalPath
		}
	} else if input.ProjectDir != "" {
		workDir = input.ProjectDir
	}

	if workDir != s.workDir {
		go s.appendRawNarrative(workDir, input.Prompt)
	}

	projectContext := readProjectContext(workDir)
	fullPrompt := input.Prompt
	if projectContext != "" {
		fullPrompt = fmt.Sprintf("Project context from %s:\n%s\n\nUser request: %s", workDir, projectContext, input.Prompt)
	}

	ctx, cancel := context.WithTimeout(r.Context(), elaborateTimeout)
	defer cancel()

	var result *elaboratedTask
	var err error

	// Try Claude first.
	result, err = s.elaborateWithClaude(ctx, workDir, fullPrompt)
	if err != nil {
		s.logger.Warn("elaborate: claude failed, falling back to gemini", "error", err)
		// Fallback to Gemini.
		result, err = s.elaborateWithGemini(ctx, workDir, fullPrompt)
		if err != nil {
			s.logger.Error("elaborate: fallback gemini also failed", "error", err)
			writeJSON(w, http.StatusBadGateway, map[string]string{
				"error": fmt.Sprintf("elaboration failed: %v", err),
			})
			return
		}
	}

	if result.Name == "" || result.Agent.Instructions == "" {
		writeJSON(w, http.StatusBadGateway, map[string]string{
			"error": "elaboration failed: missing required fields in response",
		})
		return
	}

	if result.Agent.Type == "" {
		result.Agent.Type = "claude"
	}

	sanitizeElaboratedTask(result)

	writeJSON(w, http.StatusOK, result)
}