package flow

import (
	"ai-agent/gateway"
	sessionDao "ai-agent/workflow/dao/session"
	sessionDto "ai-agent/workflow/model/dto/session"
	"ai-agent/workflow/model/entity"
	"context"
	"errors"
	"fmt"
	"time"

	"ai-agent/tools/runner"

	"gitea.redpowerfuture.com/red-future/common/tools"
	wsCommon "gitea.redpowerfuture.com/red-future/common/websocket"
	"github.com/gogf/gf/v2/frame/g"

	"github.com/gogf/gf/v2/os/glog"
	"github.com/gogf/gf/v2/util/gconv"
)

func init() {
	// 普通对话消息处理器（会话服务器 SessionWsService 见 ws_server.go）
	SessionWsService.OnMessage("agent", handleToolAgent)
	SessionWsService.OnMessage("agent_cancel", handleToolAgentCancel)
}

// 工具对话默认系统提示词
const defaultAgentSystemPrompt = "你是一个智能助手，可以调用工具完成任务。请根据任务需要选择合适的工具，参考工具返回结果，最终给出完整回答。"

// 工具对话 ReAct 最大循环步数
const defaultAgentMaxStep = 15

// handleToolAgent 处理工具对话消息：解析 payload 后异步运行 ReAct 循环，逐步推送过程事件。
// 首条消息惰性建会话（复用已存在会话），跑完后把问答/token 写入 exec_chat 落库。
func handleToolAgent(ctx context.Context, conn *wsCommon.WsConnection, payload interface{}) {
	var p sessionDto.WebSocketExecChatReq
	if err := gconv.Struct(payload, &p); err != nil {
		pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventError, Message: "参数解析失败", Error: err.Error()})
		return
	}
	if p.Question == "" {
		pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventError, Message: "内容不能为空", Error: "提问内容不能为空"})
		return
	}

	saveCtx := context.WithoutCancel(ctx)

	id := p.Id
	if g.IsEmpty(id) {
		// 会话落库：前端临时 sessionId 对应已存在会话则复用，否则新建
		err := ensureSession(saveCtx, conn.SessionId, p.Question)
		if err != nil {
			glog.Errorf(ctx, "会话创建失败: %v", err)
			pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventError, Message: "会话创建失败", Error: err.Error()})
			return
		}

		// 问答落库：
		chatId, err := sessionDao.ExecChatDao.Insert(ctx, &sessionDto.CreateExecChatReq{
			SessionId:     conn.SessionId,
			RequestParams: entity.ExecChatRequestParams{Question: p.Question},
		})
		if err != nil {
			glog.Errorf(ctx, "问答创建失败: %v", err)
			pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventError, Message: "问答创建失败", Error: err.Error()})
			return
		}
		id = chatId
		pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventRoundStart, Id: chatId})
	}

	modelTools, err := tools.Default.List(ctx)
	if err != nil {
		pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventError, Message: "工具列表获取失败", Error: err.Error()})

		errChat := recordChat(saveCtx, id, "", "工具列表获取失败", err, 0, 0, 0)
		if errChat != nil {
			glog.Errorf(ctx, "普通对话落库失败: %v", errChat)
			pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventError, Message: "对话落库失败", Error: errChat.Error()})
		}
		return
	}
	systemPrompt := p.SystemPrompt
	if systemPrompt == "" {
		systemPrompt = defaultAgentSystemPrompt
	}

	// 支持前端终止：agent 上下文可取消；落库用不带取消的 ctx（保留 request 值），保证终止后 token 仍能记录
	agentCtx, agentCancel := context.WithCancel(ctx)
	if oldCancel := getToolCancel(conn); oldCancel != nil {
		oldCancel()
	}
	conn.SetMeta("toolCancel", agentCancel)
	defer conn.SetMeta("toolCancel", nil)
	defer agentCancel()

	agent := runner.NewReActAgent(p.ModelId, conn.SessionId, modelTools, systemPrompt, defaultAgentMaxStep)
	agent.OnEvent = func(ev runner.ReActEvent) {
		pushAgentEvent(conn, ev)
	}

	start := time.Now()
	answer, runErr := agent.Run(agentCtx, p.Question)
	duration := int64(time.Since(start).Seconds())

	// 前端终止：结果/错误不推前端，仅把已产生的 token 正常落库（友好提示记「用户已终止对话」，不记原始错误）
	var errMsg string
	terminated := runErr != nil && errors.Is(runErr, context.Canceled)
	if terminated {
		errMsg = errChatTerminated.Error()
		runErr = nil
	} else if runErr != nil {
		errMsg = "对话运行失败"
		pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventError, Message: "对话运行失败", Error: runErr.Error()})
	}
	err = recordChat(saveCtx, id, answer, errMsg, runErr, agent.TotalTokens, agent.TotalCost, duration)
	if err != nil {
		glog.Errorf(ctx, "普通对话落库失败: %v", err)
		pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventError, Message: "对话落库失败", Error: err.Error()})
	}
	pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventAnswer, Answer: answer})
}

// handleToolAgentCancel 终止正在运行的对话（前端停止按钮发送 agent_cancel）
func handleToolAgentCancel(ctx context.Context, conn *wsCommon.WsConnection, _ interface{}) {
	if cancel := getToolCancel(conn); cancel != nil {
		cancel()
	}
	_ = conn.WriteJSON(&wsCommon.WsPushMsg{Type: "ack", Message: "已终止对话"})
}

// getToolCancel 获取当前 agent 运行的取消函数
func getToolCancel(conn *wsCommon.WsConnection) context.CancelFunc {
	cancel, _ := wsCommon.GetMetaT[context.CancelFunc](conn, "toolCancel")
	return cancel
}

// errChatTerminated 前端终止对话的错误标记（写入 exec_chat.error_message）
var errChatTerminated = errors.New("用户已终止对话")

// recordChat 把一次普通对话写入 exec_chat：答案传 OSS 存 result_file_url，
// 友好提示写 error_message，原始错误写 error，token 与费用（模型网关返回的累计 cost）落库
func recordChat(ctx context.Context, id int64, answer string, msg string, runErr error, totalTokens int64, totalCost float64, duration int64) error {
	var resultFileUrl string
	if runErr == nil && answer != "" {
		url, uploadErr := gateway.Upload(ctx, fmt.Sprintf("chat_%v_%d.txt", id, time.Now().UnixMilli()), []byte(answer))
		if uploadErr != nil {
			glog.Errorf(ctx, "普通对话答案上传OSS失败: %v", uploadErr)
		} else {
			resultFileUrl = url
		}
	}
	var errorDetail string
	if runErr != nil {
		errorDetail = runErr.Error()
	}
	_, err := sessionDao.ExecChatDao.Update(ctx, &sessionDto.UpdateExecChatReq{
		Id:            id,
		Duration:      duration,
		ResultFileUrl: resultFileUrl,
		TotalTokens:   int(totalTokens),
		TotalFee:      totalCost,
		ErrorMessage:  msg,
		Error:         errorDetail,
	})
	if err != nil {
		return err
	}
	// 执行成功：重新执行复用了同一条记录，OmitEmpty 的 Update 会跳过空 error_message/error，
	// 需显式清空，避免上一次失败的报错残留
	if runErr == nil {
		if _, err := sessionDao.ExecChatDao.ClearError(ctx, id); err != nil {
			glog.Errorf(ctx, "exec_chat 报错信息清空失败: %v", err)
		}
	}
	return nil
}

// pushAgentEvent 把 ReAct 过程事件转为 WS 推送消息
func pushAgentEvent(conn *wsCommon.WsConnection, ev runner.ReActEvent) {
	if conn.IsClosed() {
		return
	}
	switch ev.Type {
	case runner.ReActEventRoundStart:
		_ = conn.WriteJSON(&wsCommon.WsPushMsg{
			Type:    string(ev.Type),
			Message: "运行开始",
			Data:    map[string]interface{}{"recordId": ev.Id},
		})
	case runner.ReActEventModelCall:
		_ = conn.WriteJSON(&wsCommon.WsPushMsg{
			Type:    string(ev.Type),
			Message: "模型思考中",
			Data:    map[string]interface{}{"step": ev.Step, "maxStep": ev.MaxStep},
		})
	case runner.ReActEventToolCall:
		_ = conn.WriteJSON(&wsCommon.WsPushMsg{
			Type:    string(ev.Type),
			Message: "调用工具",
			Data:    map[string]interface{}{"description": ev.Description},
		})
	case runner.ReActEventToolResult:
		_ = conn.WriteJSON(&wsCommon.WsPushMsg{
			Type:    string(ev.Type),
			Message: "工具返回",
			Data:    map[string]interface{}{"description": ev.Description},
		})
	case runner.ReActEventAnswerChunk:
		_ = conn.WriteJSON(&wsCommon.WsPushMsg{
			Type:    string(ev.Type),
			Message: "思考中",
			Data:    map[string]interface{}{"delta": ev.Delta},
		})
	case runner.ReActEventReasoningChunk:
		_ = conn.WriteJSON(&wsCommon.WsPushMsg{
			Type:    string(ev.Type),
			Message: "思考中",
			Data:    map[string]interface{}{"delta": ev.Delta},
		})
	case runner.ReActEventAnswer:
		_ = conn.WriteJSON(&wsCommon.WsPushMsg{
			Type:    string(ev.Type),
			Message: "作答完成",
			Data:    map[string]interface{}{"answer": ev.Answer},
		})
	case runner.ReActEventError:
		_ = conn.WriteJSON(&wsCommon.WsPushMsg{Type: string(ev.Type), Message: ev.Message, Error: ev.Error})
	}
}