MuyueWorkspace/internal/workflow/workflow.go

package workflow

import (
	"encoding/json"
	"fmt"
	"strings"
)

type Phase string

const (
	PhaseIdle       Phase = "idle"
	PhaseGathering  Phase = "gathering"
	PhasePlanning   Phase = "planning"
	PhaseReviewing  Phase = "reviewing"
	PhaseExecuting  Phase = "executing"
	PhaseDone       Phase = "done"
	PhaseError      Phase = "error"
)

type Step struct {
	ID          string `json:"id"`
	Title       string `json:"title"`
	Description string `json:"description"`
	Status      string `json:"status"`
	Agent       string `json:"agent"`
	Output      string `json:"output,omitempty"`
}

type Plan struct {
	Goal         string   `json:"goal"`
	Context      string   `json:"context"`
	Questions    []string `json:"questions"`
	Answers      []string `json:"answers"`
	Steps        []Step   `json:"steps"`
	StepIndex  int      `json:"current_step"`
	PreviewFiles []PreviewFile `json:"preview_files,omitempty"`
}

type PreviewFile struct {
	Filename string `json:"filename"`
	Content  string `json:"content"`
	Type     string `json:"type"`
}

type Workflow struct {
	Phase   Phase
	Plan    *Plan
	History []string
}

func New() *Workflow {
	return &Workflow{
		Phase:   PhaseIdle,
		Plan:    &Plan{},
		History: []string{},
	}
}

func (w *Workflow) Start(goal string) {
	w.Phase = PhaseGathering
	w.Plan = &Plan{
		Goal:    goal,
		Steps:   []Step{},
		Answers: []string{},
	}
	w.History = append(w.History, fmt.Sprintf("[started] %s", goal))
}

func (w *Workflow) SetQuestions(questions []string) {
	w.Plan.Questions = questions
}

func (w *Workflow) AddAnswer(answer string) {
	w.Plan.Answers = append(w.Plan.Answers, answer)
	if len(w.Plan.Answers) >= len(w.Plan.Questions) {
		w.Phase = PhasePlanning
		w.History = append(w.History, "[gathering complete, moving to planning]")
	}
}

func (w *Workflow) SetPlan(planJSON string) error {
	var steps []Step
	if err := json.Unmarshal([]byte(planJSON), &steps); err != nil {
		if err2 := json.Unmarshal([]byte("["+planJSON+"]"), &steps); err2 != nil {
			return fmt.Errorf("parse plan: %w", err)
		}
	}
	w.Plan.Steps = steps
	w.Phase = PhaseReviewing
	w.History = append(w.History, fmt.Sprintf("[plan created] %d steps", len(steps)))
	return nil
}

func (w *Workflow) SetPreviewFiles(files []PreviewFile) {
	w.Plan.PreviewFiles = files
}

func (w *Workflow) Approve() {
	w.Phase = PhaseExecuting
	w.Plan.StepIndex = 0
	w.History = append(w.History, "[plan approved, starting execution]")
}

func (w *Workflow) Reject(feedback string) {
	w.Phase = PhasePlanning
	w.History = append(w.History, fmt.Sprintf("[plan rejected: %s]", feedback))
}

func (w *Workflow) AdvanceStep(output string) {
	if w.Plan.StepIndex < len(w.Plan.Steps) {
		w.Plan.Steps[w.Plan.StepIndex].Status = "done"
		w.Plan.Steps[w.Plan.StepIndex].Output = output
		w.Plan.StepIndex++
		w.History = append(w.History, fmt.Sprintf("[step %d done]", w.Plan.StepIndex))

		if w.Plan.StepIndex >= len(w.Plan.Steps) {
			w.Phase = PhaseDone
			w.History = append(w.History, "[all steps complete]")
		}
	}
}

func (w *Workflow) FailStep(errMsg string) {
	if w.Plan.StepIndex < len(w.Plan.Steps) {
		w.Plan.Steps[w.Plan.StepIndex].Status = "error"
		w.Plan.Steps[w.Plan.StepIndex].Output = errMsg
		w.Phase = PhaseError
		w.History = append(w.History, fmt.Sprintf("[step %d failed: %s]", w.Plan.StepIndex+1, errMsg))
	}
}

func (w *Workflow) Reset() {
	w.Phase = PhaseIdle
	w.Plan = &Plan{}
}

func (w *Workflow) CurrentStep() *Step {
	if w.Plan.StepIndex < len(w.Plan.Steps) {
		return &w.Plan.Steps[w.Plan.StepIndex]
	}
	return nil
}

func (w *Workflow) Progress() (done, total int) {
	for _, s := range w.Plan.Steps {
		if s.Status == "done" {
			done++
		}
		total++
	}
	return
}

func BuildSystemPrompt(phase Phase, plan *Plan) string {
	base := `You are muyue, an AI-powered development environment assistant.
You follow a structured workflow: GATHER requirements → PLAN → REVIEW → EXECUTE.

RULES:
- Always respond in the same language the user writes in.
- When in GATHERING phase, ask clarifying questions ONE AT A TIME to understand the requirement fully.
- When in PLANNING phase, create a detailed step-by-step plan as a JSON array of objects.
- When in REVIEWING phase, present the plan clearly and wait for approval.
- When in EXECUTING phase, execute one step at a time and report results.
- If the user wants a visual preview, generate 1-2 HTML files wrapped in a PREVIEW_JSON block.`

	switch phase {
	case PhaseGathering:
		base += fmt.Sprintf(`

CURRENT PHASE: GATHERING
Goal: %s
Questions to ask: %v
Answers received: %v
Remaining questions: %d
Ask the NEXT question that hasn't been answered yet. If all questions are answered, say "GATHERING_COMPLETE".`,
			plan.Goal, plan.Questions, plan.Answers,
			len(plan.Questions)-len(plan.Answers))

	case PhasePlanning:
		qa := ""
		for i, q := range plan.Questions {
			a := ""
			if i < len(plan.Answers) {
				a = plan.Answers[i]
			}
			qa += fmt.Sprintf("\nQ: %s\nA: %s", q, a)
		}
		base += fmt.Sprintf(`

CURRENT PHASE: PLANNING
Goal: %s
%s

Create a step-by-step plan. Output ONLY a JSON array of steps:
[
  {"id": "1", "title": "...", "description": "...", "agent": "crush|claude|muyue", "status": "pending"},
  ...
]

If the user needs a visual preview, wrap HTML in:
<<<PREVIEW_JSON>>>
[{"filename":"preview.html","content":"<html>...</html>","type":"html"}]
<<<END_PREVIEW>>>`,
			plan.Goal, qa)

	case PhaseReviewing:
		steps, _ := json.MarshalIndent(plan.Steps, "", "  ")
		base += fmt.Sprintf(`

CURRENT PHASE: REVIEWING
Present the plan below clearly and ask for approval:
%s

Say "PLAN_APPROVED" if the user approves, or "PLAN_REJECTED: <reason>" if not.`,
			string(steps))

	case PhaseExecuting:
		if plan.StepIndex < len(plan.Steps) {
			step := plan.Steps[plan.StepIndex]
			base += fmt.Sprintf(`

CURRENT PHASE: EXECUTING
Current step: %s — %s (agent: %s)
Execute this step and report the result.`,
				step.Title, step.Description, step.Agent)
		}
	}

	return base
}

func ParsePlanResponse(response string) ([]Step, error) {
	response = strings.TrimSpace(response)

	start := strings.Index(response, "[")
	end := strings.LastIndex(response, "]")
	if start == -1 || end == -1 || end <= start {
		return nil, fmt.Errorf("no JSON array found in response")
	}

	jsonStr := response[start : end+1]
	var steps []Step
	if err := json.Unmarshal([]byte(jsonStr), &steps); err != nil {
		return nil, fmt.Errorf("parse steps: %w", err)
	}

	for i := range steps {
		steps[i].Status = "pending"
	}

	return steps, nil
}

func ParsePreviewFiles(response string) []PreviewFile {
	startMarker := "<<<PREVIEW_JSON>>>"
	endMarker := "<<<END_PREVIEW>>>"
	start := strings.Index(response, startMarker)
	end := strings.Index(response, endMarker)
	if start == -1 || end == -1 {
		return nil
	}

	jsonStr := strings.TrimSpace(response[start+len(startMarker) : end])
	var files []PreviewFile
	if err := json.Unmarshal([]byte(jsonStr), &files); err != nil {
		return nil
	}
	return files
}

func ParseApproval(response string) (approved bool, feedback string) {
	lower := strings.ToLower(strings.TrimSpace(response))
	if strings.Contains(lower, "plan_approved") || strings.Contains(lower, "approved") || strings.Contains(lower, "yes") || strings.Contains(lower, "go ahead") || strings.Contains(lower, "oui") || strings.Contains(lower, "ok") {
		return true, ""
	}
	if strings.Contains(lower, "plan_rejected:") {
		parts := strings.SplitN(lower, "plan_rejected:", 2)
		if len(parts) > 1 {
			return false, strings.TrimSpace(parts[1])
		}
	}
	return false, response
}