package flow

import (
	"ai-agent/workflow/consts/node"
	flowDto "ai-agent/workflow/model/dto/flow"
	"ai-agent/workflow/model/entity"
	"context"
	"encoding/json"
	"time"

	flowDao "ai-agent/workflow/dao/flow"

	"github.com/cloudwego/eino/compose"
	"github.com/cloudwego/eino/schema"
)

// 注册 checkpoint 序列化类型
func init() {
	// ========== 1. Eino 断点核心状态（根类型） ==========
	schema.RegisterName[*flowDto.NodeExecutionState]("flow.NodeExecutionState")
	schema.RegisterName[*flowDto.NodeExecutionInput]("flow.NodeExecutionInput")
	schema.RegisterName[*flowDto.FlowExecutionInput]("flow.FlowExecutionInput")

	// ========== 2. 原有第三方类型 ==========
	schema.RegisterName[json.Number]("json.Number")
	schema.RegisterName[time.Time]("time.Time")
	schema.RegisterName[time.Duration]("time.Duration")

	// ========== 3. flowDto 内部嵌套类型 ==========
	schema.RegisterName[flowDto.ExecutedNode]("flow.ExecutedNode")

	// ========== 4. entity 核心链路类型（递归自 *entity.FlowNode） ==========
	schema.RegisterName[*entity.FlowNode]("entity.FlowNode")
	schema.RegisterName[node.NodeType]("node.NodeType")
	schema.RegisterName[*entity.SubFlowConfig]("entity.SubFlowConfig")
	//schema.RegisterName[node.NodeFormField]("node.NodeFormField")
	schema.RegisterName[entity.ModelItem]("node.ModelItem")

	// ========== 5. FlowInfo 相关（流程拓扑结构） ==========
	schema.RegisterName[entity.FlowInfo]("entity.FlowInfo")
	schema.RegisterName[entity.FlowEdge]("entity.FlowEdge")
}

// DbCheckPointStore 数据库存储实现
type DbCheckPointStore struct{}

func NewDbCheckPointStore() compose.CheckPointStore {
	return &DbCheckPointStore{}
}

func (d *DbCheckPointStore) Get(ctx context.Context, id string) ([]byte, bool, error) {
	// 去掉取消信号：断连/取消场景下 graph ctx 已取消，仅保留值，避免恢复时被阻断
	record, err := flowDao.FlowCheckpointDao.Get(context.WithoutCancel(ctx), id)
	if err != nil {
		return nil, false, err
	}
	if record == nil || record.Data == "" {
		return nil, false, nil
	}
	return []byte(record.Data), true, nil
}

func (d *DbCheckPointStore) Set(ctx context.Context, id string, val []byte) error {
	// 关键：Eino 在节点失败（Interrupt）时用 graph ctx 写 checkpoint。WS 断连/用户终止会使
	// 该 ctx 已取消，直接透传会因 "context canceled" 落库失败，断点丢失、续跑失效。
	// 去掉取消信号只保留 ctx 值，保证中断时断点必达。
	return flowDao.FlowCheckpointDao.SaveOrUpdate(context.WithoutCancel(ctx), id, string(val))
}