poimen-workflows/cmd/starter/main.go

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"go.temporal.io/sdk/client"
	"github.com/rockliang/poimen/workflows/action/llm"
	"github.com/rockliang/poimen/workflows/internal/config"
	"github.com/rockliang/poimen/workflows/internal/health"
	"github.com/rockliang/poimen/workflows/internal/logging"
	"github.com/rockliang/poimen/workflows/internal/routing"
	"github.com/rockliang/poimen/workflows/statemachine"
)

func main() {
	var (
		// Orchestrator flags
		repoPath         = flag.String("repo", "", "target repo path")
		remoteURL        = flag.String("remote", "", "remote URL")
		milestone        = flag.String("milestone", "T0", "milestone ID")
		dryRun           = flag.Bool("dry-run", false, "disable git push/merge (orchestrator) or skip submit (routing)")
		plannerModel     = flag.String("planner-model", "reasoning", "planner model ID (local-llm)")
		judgeModel       = flag.String("judge-model", "reasoning", "judge model ID (local-llm)")
		implementerModel = flag.String("implementer-model", "ornith:35b", "implementer model ID (local-llm ornith)")
		piProvider       = flag.String("pi-provider", "local-llm", "pi provider name for skills (local-llm)")
		healthCheck      = flag.Bool("health", false, "check health and exit")

		// Routing workflow flags
		routeMsg  = flag.String("route", "", "natural language message for LLM routing")
		specFile  = flag.String("spec", "", "JSON workflow spec file (direct submit, skip LLM)")
		cronSpec  = flag.Bool("cron", false, "treat spec as CronWorkflowSpec")
	)
	flag.Parse()

	// Initialize structured logging
	if err := logging.InitLogger(); err != nil {
		log.Fatalf("failed to initialize logger: %v", err)
	}
	defer logging.Sync()

	// Load configuration first
	cfg, err := config.LoadConfig()
	if err != nil {
		logging.Fatal("failed to load config", logging.Err(err))
	}

	// Connect to Temporal
	logging.Info("connecting to Temporal", logging.String("hostPort", cfg.Temporal.HostPort), logging.String("namespace", cfg.Temporal.Namespace))
	c, err := client.Dial(client.Options{
		HostPort:  cfg.Temporal.HostPort,
		Namespace: cfg.Temporal.Namespace,
	})
	if err != nil {
		logging.Fatal("failed to connect to temporal", logging.Err(err))
	}
	defer c.Close()

	// If health check requested, do it and exit
	if *healthCheck {
		logging.Info("running health check")
		healthChecker := health.NewChecker(c)
		report := healthChecker.Check(context.Background())
		jsonReport, _ := report.ToJSON()
		fmt.Println(string(jsonReport))
		if report.Status != health.StatusHealthy {
			logging.Fatal("health check failed")
		}
		return
	}

	// Handle routing workflow mode
	if *routeMsg != "" || *specFile != "" {
		runRoutingWorkflow(c, *routeMsg, *specFile, *cronSpec, *dryRun)
		return
	}

	// Validate required flags for orchestrator workflow
	if *repoPath == "" || *remoteURL == "" {
		logging.Fatal("--repo and --remote flags are required (or use --route/--spec for routing workflow)")
	}



	// Build OrchestratorInput
	input := statemachine.OrchestratorInput{
		TargetRepoPath:     *repoPath,
		RemoteURL:          *remoteURL,
		Milestone:          *milestone,
		DryRun:             *dryRun,
		MaxCyclesBeforeCAN: 100,
		PiProvider:         *piProvider,
		Config: statemachine.OrchestratorConfig{
			SystemPrompt: "You are an expert software developer orchestrating multi-agent work.",
			Skills:       []statemachine.SkillRef{},
			RolePrompts: map[string]statemachine.PromptSpec{
				"planner": {
					TemplateRef: "planner/default.tmpl",
					Model: statemachine.ModelSpec{
						ModelID:  *plannerModel,
						Thinking: "adaptive",
						Effort:   "high",
					},
				},
				"judge": {
					TemplateRef: "judge/default.tmpl",
					Model: statemachine.ModelSpec{
						ModelID:  *judgeModel,
						Thinking: "adaptive",
						Effort:   "high",
					},
				},
				"implementer": {
					TemplateRef: "implementer/default.tmpl",
					Model: statemachine.ModelSpec{
						ModelID: *implementerModel,
					},
				},
			},
			Tuning: statemachine.NewActivityTuning(),
		},
	}

	// Health check: verify local LLM API is reachable
	logging.Info("checking local LLM API connectivity", logging.String("url", "https://api.riotpiao.com"))
	llmClient, err := llm.NewClient()
	if err != nil {
		logging.Fatal("failed to create LLM client", logging.Err(err))
	}
	if err := llmClient.HealthCheck(context.Background()); err != nil {
		logging.Fatal("local LLM API health check failed", logging.Err(err), logging.String("hint", "ensure homelab-frontend gateway is running and accessible"))
	}
	logging.Info("local LLM API is reachable", logging.String("planner-model", *plannerModel), logging.String("judge-model", *judgeModel), logging.String("implementer-model", *implementerModel))

	// Start workflow
	workflowID := "orch-" + strings.ReplaceAll(*repoPath, "/", "-")
	logging.Info("starting orchestrator workflow", logging.String("workflowID", workflowID), logging.String("repo", *repoPath))
	run, err := c.ExecuteWorkflow(context.Background(), client.StartWorkflowOptions{
		ID:        workflowID,
		TaskQueue: "poimen-taskqueue",
	}, statemachine.OrchestratorWorkflow, input)
	if err != nil {
		logging.Fatal("failed to start workflow", logging.Err(err))
	}

	fmt.Printf("\n=== Workflow Started ===\n")
	fmt.Printf("Workflow ID: %s\n", workflowID)
	fmt.Printf("Task Queue:  poimen-taskqueue\n")
	fmt.Printf("\n=== Model Configuration ===\n")
	fmt.Printf("Planner Model:     %s\n", *plannerModel)
	fmt.Printf("Judge Model:       %s\n", *judgeModel)
	fmt.Printf("Implementer Model: %s\n", *implementerModel)
	fmt.Printf("\n=== Monitoring ===\n")
	fmt.Printf("Web UI: http://%s:8080/namespaces/%s/workflows/%s\n",
		strings.Split(cfg.Temporal.HostPort, ":")[0], cfg.Temporal.Namespace, workflowID)

	// Optionally wait for completion (with timeout)
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute)
	defer cancel()

	var result statemachine.OrchestratorOutput
	if err := run.Get(ctx, &result); err != nil {
		fmt.Printf("\nWorkflow initiated (execution in progress).\n")
		fmt.Printf("Check the Web UI for real-time status updates.\n")
	} else {
		fmt.Printf("\nWorkflow completed: %+v\n", result)
	}
}

// runRoutingWorkflow handles --route and --spec flags
func runRoutingWorkflow(c client.Client, routeMsg, specFile string, isCron, dryRun bool) {
	ctx := context.Background()

	var spec *routing.WorkflowSpec
	var cronSpec *routing.CronWorkflowSpec

	if specFile != "" {
		// Load spec from file
		data, err := os.ReadFile(specFile)
		if err != nil {
			logging.Fatal("failed to read spec file", logging.Err(err))
		}

		validator := routing.NewValidator(nil) // nil KB = skip activity validation
		if isCron {
			cronSpec = &routing.CronWorkflowSpec{}
			if err := json.Unmarshal(data, cronSpec); err != nil {
				logging.Fatal("failed to parse cron spec", logging.Err(err))
			}
			// Validate
			result := validator.ValidateCronWorkflowSpec(cronSpec)
			if !result.Valid {
				logging.Fatal("invalid cron spec", logging.String("errors", result.String()))
			}
		} else {
			spec = &routing.WorkflowSpec{}
			if err := json.Unmarshal(data, spec); err != nil {
				logging.Fatal("failed to parse spec", logging.Err(err))
			}
			// Validate
			result := validator.ValidateWorkflowSpec(spec)
			if !result.Valid {
				logging.Fatal("invalid spec", logging.String("errors", result.String()))
			}
		}
	} else {
		// Use LLM router
		logging.Info("routing message via LLM", logging.String("message", routeMsg))

		kb, err := routing.LoadKnowledgeBaseFromDefaultPath()
		if err != nil {
			logging.Fatal("failed to load knowledge base", logging.Err(err))
		}

		router, err := routing.NewLLMRouterDefault(kb)
		if err != nil {
			logging.Fatal("failed to create LLM router", logging.Err(err))
		}

		output, err := router.Route(ctx, routing.LLMRouterInput{Message: routeMsg})
		if err != nil {
			logging.Fatal("LLM routing failed", logging.Err(err))
		}

		if output.IsCron {
			cronSpec = output.CronSpec
			fmt.Printf("\n=== Generated Cron Spec ===\n")
			fmt.Printf("Name: %s\n", cronSpec.Name)
			fmt.Printf("Schedule: %s\n", cronSpec.Schedule)
			fmt.Printf("States: %d\n", len(cronSpec.States))
		} else {
			spec = output.Spec
			fmt.Printf("\n=== Generated Workflow Spec ===\n")
			fmt.Printf("Name: %s\n", spec.Name)
			fmt.Printf("States: %d\n", len(spec.States))
		}
	}

	if dryRun {
		fmt.Printf("\n[dry-run] Spec generated but not submitted\n")
		if spec != nil {
			data, _ := json.MarshalIndent(spec, "", "  ")
			fmt.Printf("%s\n", data)
		} else if cronSpec != nil {
			data, _ := json.MarshalIndent(cronSpec, "", "  ")
			fmt.Printf("%s\n", data)
		}
		return
	}

	// Submit to Temporal
	if cronSpec != nil {
		// For cron, we'd use Temporal's schedule feature
		// For now, just start as regular workflow (cron scheduling TBD)
		spec = &routing.WorkflowSpec{
			Name:   cronSpec.Name,
			Input:  cronSpec.Input,
			States: cronSpec.States,
		}
		logging.Warn("cron scheduling not yet implemented, running as one-shot workflow")
	}

	workflowID := "routing-" + spec.Name + "-" + time.Now().Format("20060102-150405")
	input := statemachine.RoutingWorkflowInput{Spec: spec}

	run, err := c.ExecuteWorkflow(ctx, client.StartWorkflowOptions{
		ID:        workflowID,
		TaskQueue: "poimen-taskqueue",
	}, statemachine.RoutingWorkflow, input)
	if err != nil {
		logging.Fatal("failed to start routing workflow", logging.Err(err))
	}

	fmt.Printf("\n=== Routing Workflow Started ===\n")
	fmt.Printf("Workflow ID: %s\n", workflowID)
	fmt.Printf("Run ID: %s\n", run.GetRunID())

	// Wait briefly for result
	waitCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()

	var result statemachine.RoutingWorkflowOutput
	if err := run.Get(waitCtx, &result); err != nil {
		fmt.Printf("\nWorkflow running (check Temporal UI for status)\n")
	} else {
		fmt.Printf("\nWorkflow completed: %s\n", result.Status)
		if len(result.StepResults) > 0 {
			for step, res := range result.StepResults {
				fmt.Printf("  %s: %v\n", step, res)
			}
		}
	}
}