package flow

import (
	"context"
	"encoding/json"
	"errors"
	"time"

	"github.com/gogf/gf/v2/frame/g"

	"ai-agent/gateway"
	"ai-agent/workflow/consts/model"
	flowDao "ai-agent/workflow/dao/flow"
	"ai-agent/workflow/model/entity"
)

// asyncRecoverWaitTimeout in-flight 行重订阅等结果的超时上限。
// 任务可能仍执行中（消息未发布）或提交即失败（不会发布消息）；超时视为结果未知，清记录重提。
const asyncRecoverWaitTimeout = 10 * time.Minute

type asyncAction int

const (
	asyncActionResubmit asyncAction = iota // 无记录 / failed / done 空结果 → 重新提交
	asyncActionReuse                       // done 有结果 → 复用，不重调
	asyncActionFinalize                    // in-flight → 重订阅 msgTopic 收尾
)

// asyncCallAction 异步任务缓存行动决策（纯函数，可单测）
func asyncCallAction(rec *entity.FlowAsyncTask) asyncAction {
	if rec == nil {
		return asyncActionResubmit
	}
	switch rec.State {
	case flowDao.FlowAsyncStateDone:
		if rec.Result != "" && rec.Result != "{}" {
			return asyncActionReuse
		}
		return asyncActionResubmit
	case flowDao.FlowAsyncStateFailed:
		return asyncActionResubmit
	default: // in-flight
		return asyncActionFinalize
	}
}

// AsyncModelCallWithRecovery 统一异步模型调用入口（spec §5.4）：
// 提交时把 task_id/msg_topic 落库 flow_async_task，崩溃后重订阅 msg_topic 拿回已完成结果复用，不重复调用。
// 同步模型直接走 gateway.ModelCallResult，不落库（无恢复语义）。
func AsyncModelCallWithRecovery(ctx context.Context, execId int64, nodeId string, segIdx int, modelId int64, responseType model.ResponseType, sessionId string, requestParams map[string]any, businessParams map[string]any) (*gateway.ModelCallRes, error) {
	if responseType == nil || *responseType != *model.ResponseTypeAsync.Code() {
		return gateway.ModelCallResult(ctx, modelId, responseType, sessionId, requestParams, businessParams)
	}

	rec, err := flowDao.FlowAsyncTaskDao.Get(ctx, execId, nodeId, segIdx)
	if err != nil {
		return nil, err
	}
	switch asyncCallAction(rec) {
	case asyncActionReuse:
		return unmarshalModelCallRes(rec.Result)
	case asyncActionFinalize:
		// 重订阅 msgTopic 收尾：拿回已发布结果（消息在 JetStream 保留 7 天）；
		// 成功 → 落 done 复用；超时/订阅失败 → 清记录重提
		waitCtx, cancel := context.WithTimeout(ctx, asyncRecoverWaitTimeout)
		res, waitErr := gateway.WaitModelCallResult(waitCtx, rec.MsgTopic)
		cancel()
		if waitErr != nil {
			if errors.Is(waitErr, context.DeadlineExceeded) || errors.Is(waitErr, context.Canceled) {
				_ = flowDao.FlowAsyncTaskDao.DeleteByKey(ctx, execId, nodeId, segIdx)
				break // 落入下方重新提交
			}
			return nil, waitErr
		}
		_ = flowDao.FlowAsyncTaskDao.UpdateByKey(ctx, execId, nodeId, segIdx, flowDao.FlowAsyncStateDone, marshalModelCallRes(res))
		return res, nil
	case asyncActionResubmit:
		// 清残留（failed 或 done 空结果）
		if rec != nil {
			_ = flowDao.FlowAsyncTaskDao.DeleteByKey(ctx, execId, nodeId, segIdx)
		}
	}

	// 重新提交：先落库 in-flight（task_id/msg_topic），等待结果
	res, msgTopic, err := gateway.SubmitModelCall(ctx, modelId, responseType, sessionId, requestParams, businessParams)
	if err != nil {
		return nil, err
	}
	if err := flowDao.FlowAsyncTaskDao.Upsert(ctx, execId, nodeId, segIdx, modelId, res.TaskId, msgTopic); err != nil {
		return nil, err
	}
	waitRes, waitErr := gateway.WaitModelCallResult(ctx, msgTopic)
	if waitErr != nil {
		// 结果失败（model error/取消）：落 failed，调用方（段重试/恢复）决定后续
		_ = flowDao.FlowAsyncTaskDao.UpdateByKey(ctx, execId, nodeId, segIdx, flowDao.FlowAsyncStateFailed, "")
		return nil, waitErr
	}
	_ = flowDao.FlowAsyncTaskDao.UpdateByKey(ctx, execId, nodeId, segIdx, flowDao.FlowAsyncStateDone, marshalModelCallRes(waitRes))
	return waitRes, nil
}

func marshalModelCallRes(res *gateway.ModelCallRes) string {
	b, err := json.Marshal(res)
	if err != nil {
		g.Log().Warningf(context.Background(), "序列化模型调用结果失败: %v", err)
		return "{}"
	}
	return string(b)
}

func unmarshalModelCallRes(s string) (*gateway.ModelCallRes, error) {
	res := new(gateway.ModelCallRes)
	if err := json.Unmarshal([]byte(s), res); err != nil {
		return nil, err
	}
	return res, nil
}