// Async Ask worker: start on one process, poll from another.
//
// Documented at https://appcd-dev.github.io/stackgen-sdk/code-samples/async-ask-worker/
//
//	export STACKGEN_URL=https://app.stackgen.com
//	export STACKGEN_TOKEN=stackgen_…
//	export STACKGEN_WORKSPACE_ID=<workspace-uuid>
//	go run ./examples/go/recipes/ask-async-worker.go --start-only
//	go run ./examples/go/recipes/ask-async-worker.go --follow-up \
//	  --session-id <uuid> --trace-id <id>
package main

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

	stackgen "github.com/appcd-dev/stackgen-sdk/clients/go"
	"github.com/appcd-dev/stackgen-sdk/clients/go/aiden"
)

func envOr(keys ...string) string {
	for _, key := range keys {
		if v := strings.TrimSpace(os.Getenv(key)); v != "" {
			return v
		}
	}
	return ""
}

func printJSON(v any) {
	raw, _ := json.MarshalIndent(v, "", "  ")
	fmt.Println(string(raw))
}

func main() {
	workflow := flag.String("workflow", envOr("STACKGEN_WORKFLOW"), "Workflow name pinned on start")
	message := flag.String("message", "Investigate the Acme checkout CPU alert in prod", "Ask message")
	correlationID := flag.String("correlation-id", "acme-corr-123", "Sample correlation id")
	incidentID := flag.String("incident-id", "ACME-1001", "Sample incident id")
	startOnly := flag.Bool("start-only", false, "Fire-and-forget start; print ids and exit")
	followUp := flag.Bool("follow-up", false, "One-shot GetRunStatus for a previously started run")
	sessionID := flag.String("session-id", "", "Session id from a previous start")
	traceID := flag.String("trace-id", "", "Trace id from a previous start")
	readyArtifact := flag.String("ready-artifact", "session-report.md", "Artifact that means the run is consumable")
	downloadDir := flag.String("download-dir", "", "When following up and ready, download the artifact here")
	timeoutSeconds := flag.Int("timeout-seconds", 1800, "Client timeout")
	flag.Parse()

	if strings.TrimSpace(*workflow) == "" {
		*workflow = "acme-incident-triage"
	}
	client, err := newClient(*timeoutSeconds, 5)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	ctx := context.Background()
	if *followUp {
		os.Exit(follow(ctx, client, *sessionID, *traceID, *readyArtifact, *downloadDir))
	}
	_ = startOnly
	os.Exit(start(ctx, client, *workflow, *message, *correlationID, *incidentID))
}

func newClient(timeoutSeconds, poll int) (*stackgen.Client, error) {
	base := envOr("STACKGEN_URL", "STACKGEN_BASE_URL")
	token := envOr("STACKGEN_TOKEN", "STACKGEN_API_TOKEN")
	org := envOr("STACKGEN_WORKSPACE_ID", "STACKGEN_PROJECT", "STACKGEN_PROJECT_ID", "STACKGEN_ORG_ID")
	if base == "" || token == "" || org == "" {
		return nil, fmt.Errorf("set STACKGEN_URL, STACKGEN_TOKEN, and STACKGEN_WORKSPACE_ID")
	}
	return stackgen.NewClient(stackgen.Config{
		BaseURL:             base,
		APIToken:            token,
		WorkspaceID:         org,
		TimeoutSeconds:      timeoutSeconds,
		PollIntervalSeconds: poll,
	})
}

func start(ctx context.Context, client *stackgen.Client, workflow, message, correlationID, incidentID string) int {
	started, err := client.Aiden.StartAsk(ctx, message, aiden.StartAskOptions{
		EntityRefs: []aiden.EntityRef{{Name: workflow, Type: "workflow"}},
		SourceApp:  "acme-automation",
		SourceMetadata: map[string]string{
			"correlation_id": correlationID,
			"incident_id":    incidentID,
		},
	})
	if err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		return 1
	}
	payload := map[string]any{
		"session_id":     started.SessionID,
		"trace_id":       started.TraceID,
		"run_id":         started.RunID,
		"workflow":       workflow,
		"correlation_id": correlationID,
		"incident_id":    incidentID,
	}
	printJSON(payload)
	if started.SessionID == "" && started.TraceID == "" {
		fmt.Fprintln(os.Stderr, "error: start missing session_id and trace_id")
		return 1
	}
	return 0
}

func follow(ctx context.Context, client *stackgen.Client, sessionID, traceID, readyArtifact, downloadDir string) int {
	if strings.TrimSpace(sessionID) == "" && strings.TrimSpace(traceID) == "" {
		fmt.Fprintln(os.Stderr, "error: --follow-up requires --session-id and/or --trace-id")
		return 1
	}
	status, err := client.Aiden.GetRunStatus(ctx, aiden.RunStatusQuery{
		SessionID:     sessionID,
		TraceID:       traceID,
		ReadyArtifact: readyArtifact,
	})
	if err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		return 1
	}
	printJSON(status.ToMap())
	if status.Phase == aiden.RunPhaseError {
		return 1
	}
	readyOK := true
	if readyArtifact != "" {
		readyOK = status.Ready
	} else if status.Phase == aiden.RunPhasePending || status.Phase == aiden.RunPhaseRunning || status.Phase == aiden.RunPhaseUnknown {
		readyOK = false
	}
	if readyOK && downloadDir != "" && status.SessionID != "" && (readyArtifact == "" || status.Ready) {
		name := readyArtifact
		if name == "" {
			name = "session-report.md"
		}
		found := false
		for _, artifact := range status.ArtifactNames {
			if artifact == name {
				found = true
				break
			}
		}
		if found {
			if err := os.MkdirAll(downloadDir, 0o755); err != nil {
				fmt.Fprintln(os.Stderr, "error:", err)
				return 1
			}
			dest := downloadDir + "/" + name
			if _, err := client.Aiden.DownloadSessionArtifact(ctx, status.SessionID, name, dest); err != nil {
				fmt.Fprintln(os.Stderr, "error:", err)
				return 1
			}
			fmt.Fprintf(os.Stderr, "downloaded %s\n", dest)
		}
	}
	if readyOK {
		return 0
	}
	return 2
}