package session

import (
	"context"
	"os"
	"testing"
	"time"

	"ai-agent/workflow/consts/flow"
	sessionDto "ai-agent/workflow/model/dto/session"
	"ai-agent/workflow/model/entity"

	"gitea.redpowerfuture.com/red-future/common/beans"
	_ "github.com/gogf/gf/contrib/drivers/pgsql/v2"
)

// TestExecWorkflowRetryColsRoundTrip 验证三新列读写 + ListRecoverable 三类判定（需 AI_AGENT_TEST_DB=1 且本地 PG 可用）
func TestExecWorkflowRetryColsRoundTrip(t *testing.T) {
	if !isTestDBEnabled(t) {
		t.Skip("skip: AI_AGENT_TEST_DB not set")
	}
	// gfdb 的 Hook 依赖 ctx 中注入用户信息（租户/创建人），否则 Insert 报 "token 数据为空"
	ctx := context.WithValue(context.Background(), "user", &beans.User{
		UserName: "test-retry", TenantId: 1,
	})
	id, err := ExecWorkflowDao.Insert(ctx, &sessionDto.CreateWorkflowReq{
		SessionId: "test-retry", FlowId: 1, NodeGroupId: "ng",
		Status:        flow.FlowExecutionStatusRunning.Code(),
		LastHeartbeat: time.Now().UnixMilli(),
	})
	if err != nil {
		t.Fatalf("insert: %v", err)
	}
	defer ExecWorkflowDao.Delete(ctx, &sessionDto.DeleteExecWorkflowReq{Id: []int64{id}})

	if err := ExecWorkflowDao.TouchHeartbeat(ctx, id); err != nil {
		t.Fatalf("touch heartbeat: %v", err)
	}
	if err := ExecWorkflowDao.UpdateRetry(ctx, id, 1, 1); err != nil {
		t.Fatalf("update retry: %v", err)
	}
	got, err := ExecWorkflowDao.GetById(ctx, id)
	if err != nil || got == nil {
		t.Fatalf("get: %v", err)
	}
	if got.Retryable != 1 || got.RetryCount != 1 || got.LastHeartbeat == 0 {
		t.Fatalf("cols not persisted: retryable=%d retryCount=%d heartbeat=%d", got.Retryable, got.RetryCount, got.LastHeartbeat)
	}

	now := time.Now().UnixMilli()
	stale := now - int64(120*time.Second/time.Millisecond)
	if err := ExecWorkflowDao.UpdateRetry(ctx, id, 0, 0); err != nil {
		t.Fatalf("update retry reset: %v", err)
	}
	// 情况1：status=3 retryable=1 retry_count<N → 可恢复
	_, err = ExecWorkflowDao.Update(ctx, &sessionDto.UpdateWorkflowReq{Id: id, Status: flow.FlowExecutionStatusFailed.Code()})
	if err != nil {
		t.Fatalf("update failed: %v", err)
	}
	if err := ExecWorkflowDao.UpdateRetry(ctx, id, 1, 0); err != nil {
		t.Fatalf("update retry: %v", err)
	}
	rows, err := ExecWorkflowDao.ListRecoverable(ctx, now, stale, 2)
	if err != nil {
		t.Fatalf("list recoverable: %v", err)
	}
	if !containsID(rows, id) {
		t.Fatalf("status=3 retryable=1 应可恢复，未命中 id=%d", id)
	}
	// 情况2：retry_count 耗尽 → 不可恢复
	if err := ExecWorkflowDao.UpdateRetry(ctx, id, 1, 2); err != nil {
		t.Fatalf("update retry exhaust: %v", err)
	}
	rows, err = ExecWorkflowDao.ListRecoverable(ctx, now, stale, 2)
	if err != nil {
		t.Fatalf("list recoverable: %v", err)
	}
	if containsID(rows, id) {
		t.Fatalf("retry 耗尽不应可恢复")
	}
	// 情况3：status=3 retryable=0（用户取消）→ 不可恢复
	if err := ExecWorkflowDao.UpdateRetry(ctx, id, 0, 0); err != nil {
		t.Fatalf("update retry: %v", err)
	}
	rows, err = ExecWorkflowDao.ListRecoverable(ctx, now, stale, 2)
	if err != nil {
		t.Fatalf("list recoverable: %v", err)
	}
	if containsID(rows, id) {
		t.Fatalf("用户取消不可恢复")
	}
}

func containsID(rows []*entity.ExecWorkflow, id int64) bool {
	for _, r := range rows {
		if r.Id == id {
			return true
		}
	}
	return false
}

// isTestDBEnabled 测试辅助：AI_AGENT_TEST_DB=1 时连接真实 PG
func isTestDBEnabled(t *testing.T) bool {
	return os.Getenv("AI_AGENT_TEST_DB") == "1"
}