// ❌ 第一版作法：直接把幾 MB 的原始圖片 []byte 塞進 Session State
InitialState: map[string]any{
    StateKeyInputType:  RouteIngest,
    StateKeyImageBytes: imageBytes, // 3MB JPEG 在 JSON 序列化後膨脹為 ~4MB Base64 字串！
}

// ❌ 第一版作法：只依賴單機記憶體 map 判斷是否有暫停中的 HITL
type Engine struct {
    mu             sync.RWMutex
    pendingSession map[string]pendingHITL // userID -> {SessionID, InterruptID}
}

package sessions

import (
	"context"
	"log/slog"
	"os"
	"strings"

	"google.golang.org/adk/v2/artifact"
	"google.golang.org/adk/v2/artifact/gcsartifact"
)

// NewArtifactService returns an ADK artifact.Service backed by Google Cloud Storage
// when ARTIFACT_BUCKET_NAME or LOGS_BUCKET_NAME is configured, and falls back to
// artifact.InMemoryService() for local development and unit tests.
func NewArtifactService(ctx context.Context) (artifact.Service, error) {
	bucketName := resolveArtifactBucketName()
	if bucketName == "" {
		slog.Info("Using in-memory artifact service (set ARTIFACT_BUCKET_NAME or LOGS_BUCKET_NAME for GCS)")
		return artifact.InMemoryService(), nil
	}

	svc, err := gcsartifact.NewService(ctx, bucketName)
	if err != nil {
		return nil, err
	}
	slog.Info("Using GCS artifact service", "bucket", bucketName)
	return svc, nil
}

// main.go — 同時將 SessionService 與 ArtifactService 注入 ADK Universal Launcher
config := &launcher.Config{
    SessionService:  sessionService,
    ArtifactService: artifactService,
    AgentLoader:     NewSingleAgentLoader(rootAgent),
    A2AOptions:      a2aOpts,
}

const (
	StateKeyImageBytes        = "image_bytes"         // 僅保留向下相容 fallback
	StateKeyCardImageArtifact = "card_image_artifact" // Artifact 檔案名稱
	StateKeyCardImageVersion  = "card_image_version"  // Artifact 版本號 (int)

	DefaultCardImageArtifactName = "business_card.jpg"
)

initialState := map[string]any{
	StateKeyInputType:    RouteIngest,
	StateKeyRecentEvents: recentEvents,
}

if e.artifactService != nil && len(imageBytes) > 0 {
	part := genai.NewPartFromBytes(imageBytes, detectImageMIMEType(imageBytes))
	saveResp, saveErr := e.artifactService.Save(ctx, &artifact.SaveRequest{
		AppName:   appName,
		UserID:    userID,
		SessionID: sess.ID(),
		FileName:  DefaultCardImageArtifactName,
		Part:      part,
	})
	if saveErr == nil && saveResp != nil {
		initialState[StateKeyCardImageArtifact] = DefaultCardImageArtifactName
		initialState[StateKeyCardImageVersion] = saveResp.Version
	}
}

// internal/workflow/nodes_ingest.go — RouterNode
if getStringState(ctx, StateKeyCardImageArtifact) == "" && len(getBytesState(ctx, StateKeyImageBytes)) == 0 && ctx.UserContent() != nil {
	for _, part := range ctx.UserContent().Parts {
		if part != nil && part.InlineData != nil && len(part.InlineData.Data) > 0 {
			if ctx.Artifacts() != nil {
				mimeType := part.InlineData.MIMEType
				if mimeType == "" {
					mimeType = detectImageMIMEType(part.InlineData.Data)
				}
				artifactPart := genai.NewPartFromBytes(part.InlineData.Data, mimeType)
				if saveResp, saveErr := ctx.Artifacts().Save(ctx, DefaultCardImageArtifactName, artifactPart); saveErr == nil {
					setState(ctx, ev, StateKeyCardImageArtifact, DefaultCardImageArtifactName)
					setState(ctx, ev, StateKeyCardImageVersion, saveResp.Version)
				} else {
					setState(ctx, ev, StateKeyImageBytes, part.InlineData.Data)
				}
			}
			if getStringState(ctx, StateKeyInputType) == "" {
				setState(ctx, ev, StateKeyInputType, RouteIngest)
			}
			break
		}
	}
}

func loadCardImageBytes(ctx agent.Context) []byte {
	if artifactName := getStringState(ctx, StateKeyCardImageArtifact); artifactName != "" && ctx.Artifacts() != nil {
		version := getIntState(ctx, StateKeyCardImageVersion)
		var (
			resp *artifact.LoadResponse
			err  error
		)
		if version > 0 {
			resp, err = ctx.Artifacts().LoadVersion(ctx, artifactName, version)
		} else {
			resp, err = ctx.Artifacts().Load(ctx, artifactName)
		}
		if err == nil && resp != nil && resp.Part != nil && resp.Part.InlineData != nil {
			return resp.Part.InlineData.Data
		}
	}
	return getBytesState(ctx, StateKeyImageBytes) // Fallback
}

// resolvePendingHITL checks the local cache first and falls back to querying
// the shared session.Service for any session belonging to userID with a non-empty
// StateKeyActiveInterruptID. This enables stateless cross-instance HITL resume
// on Cloud Run / Vertex AI Agent Runtime.
func (e *Engine) resolvePendingHITL(ctx context.Context, userID string) (pendingHITL, bool) {
	// 1. Fast path：若同一個實例記憶體中已有紀錄，先向 sessionService 驗證是否仍處於暫停狀態
	e.mu.RLock()
	cached, ok := e.pendingSession[userID]
	e.mu.RUnlock()
	if ok && cached.SessionID != "" && cached.InterruptID != "" {
		if getResp, err := e.sessionService.Get(ctx, &session.GetRequest{
			AppName:   appName,
			UserID:    userID,
			SessionID: cached.SessionID,
		}); err == nil && getResp != nil && getResp.Session != nil {
			if activeID, err := getResp.Session.State().Get(StateKeyActiveInterruptID); err == nil {
				if s, ok := activeID.(string); ok && strings.TrimSpace(s) != "" {
					cached.InterruptID = strings.TrimSpace(s)
					return cached, true
				}
			}
		}
	}

	// 2. Stateless Slow Path：當落在全新容器實例時，直接查詢共享的 sessionService！
	listResp, err := e.sessionService.List(ctx, &session.ListRequest{
		AppName: appName,
		UserID:  userID,
	})
	if err != nil || listResp == nil {
		return pendingHITL{}, false
	}

	var (
		matched     pendingHITL
		latestTime  time.Time
		foundActive bool
	)
	for _, sess := range listResp.Sessions {
		if sess == nil || sess.State() == nil {
			continue
		}
		val, err := sess.State().Get(StateKeyActiveInterruptID)
		if err != nil || val == nil {
			continue
		}
		interruptID, ok := val.(string)
		interruptID = strings.TrimSpace(interruptID)
		if !ok || interruptID == "" {
			continue
		}
		updatedAt := sess.LastUpdateTime()
		if !foundActive || updatedAt.After(latestTime) {
			foundActive = true
			latestTime = updatedAt
			matched = pendingHITL{
				SessionID:   sess.ID(),
				InterruptID: interruptID,
			}
		}
	}

	if foundActive {
		e.mu.Lock()
		e.pendingSession[userID] = matched
		e.mu.Unlock()
		return matched, true
	}
	return pendingHITL{}, false
}

func TestCrossInstanceHITLResumeAndArtifactStorage(t *testing.T) {
	ctx := context.Background()
	cardStore := store.NewMemoryStore()
	mockLLM := NewMockLLM()

	// 模擬雲端共享的 Agent Platform Session Service 與 GCS Artifact Service
	sharedSessionSvc := session.InMemoryService()
	sharedArtifactSvc := artifact.InMemoryService()

	// 建立 3 個完全獨立的 Engine 實例（各自擁有獨立且空白的 pendingSession map）
	instanceA, _ := NewEngine(cardStore, mockLLM,
		WithSessionService(sharedSessionSvc),
		WithArtifactService(sharedArtifactSvc),
	)
	instanceB, _ := NewEngine(cardStore, mockLLM,
		WithSessionService(sharedSessionSvc),
		WithArtifactService(sharedArtifactSvc),
	)
	instanceC, _ := NewEngine(cardStore, mockLLM,
		WithSessionService(sharedSessionSvc),
		WithArtifactService(sharedArtifactSvc),
	)

	userID := "U_cross_instance_user"
	fakeImage := []byte{0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x10, 'J', 'F', 'I', 'F'}

	// Turn 1 (落在 Instance A)：上傳名片照片 -> 寫入 Artifact 並暫停於 HITLConfirmNode
	res1, _ := instanceA.StartIngest(ctx, userID, fakeImage, []string{"在 Cloud Summit 認識"})
	// 驗證 Session State 中沒有肥大的 raw image_bytes，而是輕量的 card_image_artifact="business_card.jpg"

	// Turn 2 (落在全新的 Instance B)：使用者輸入修改指令「職稱改為技術長」
	// Instance B 自動從 sharedSessionSvc 探測到 Instance A 留下的 active_interrupt_id！
	res2, _ := instanceB.ResumeHITL(ctx, userID, "職稱改為技術長")
	// 驗證職稱已更新為「技術長」，且再次進入 HITL 暫停等待確認

	// Turn 3 (落在全新的 Instance C)：使用者點擊「✅ 確認存檔」
	res3, _ := instanceC.ResumeHITL(ctx, userID, "action=confirm_save")
	// 驗證名片成功寫入 CardStore、產出 AI 破冰信，且所有實例的 HasPendingHITL 皆歸零！
}