package eino

import (
	"context"
	"encoding/base64"
	"io"
	"time"

	"github.com/cloudwego/eino/compose"

	"github.com/hhs/camtalk/internal/logger"
	"github.com/hhs/camtalk/internal/models"
	"github.com/hhs/camtalk/internal/orchestrator"
	"github.com/hhs/camtalk/internal/session"
)

// ctxKeySessionID sessionID 的 context key。
type ctxKeySessionID struct{}

// WithSessionID 将 sessionID 注入 context。
func WithSessionID(ctx context.Context, sessionID string) context.Context {
	return context.WithValue(ctx, ctxKeySessionID{}, sessionID)
}

// EinoOrchestrator 实现 orchestrator.Orchestrator 接口。
// 将 Eino Graph 包装为现有接口，WS Handler 几乎不用改。
type EinoOrchestrator struct {
	graph      *PipelineGraph
	sessionMgr session.Manager
	model      string
	callbacks  compose.Option // 运行时 Callback option
}

// NewEinoOrchestrator 创建 Eino 编排器适配器。
func NewEinoOrchestrator(graph *PipelineGraph, sessionMgr session.Manager, model string) *EinoOrchestrator {
	return &EinoOrchestrator{
		graph:      graph,
		sessionMgr: sessionMgr,
		model:      model,
		callbacks:  compose.WithCallbacks(BuildCallbackHandler()),
	}
}

// ProcessQuery 实现 orchestrator.Orchestrator 接口。
func (e *EinoOrchestrator) ProcessQuery(
	ctx context.Context,
	sessionID string,
	req models.WsQuery,
	sender orchestrator.Sender,
) error {
	log := logger.Log
	startTime := time.Now()

	// 1. 设置活跃请求
	if err := e.sessionMgr.SetActiveRequest(ctx, sessionID, req.RequestID); err != nil {
		log.Errorw("设置活跃请求失败", "error", err)
	}
	defer e.sessionMgr.ClearActiveRequest(ctx, sessionID)

	// 2. 获取会话配置
	sess, err := e.sessionMgr.Get(ctx, sessionID)
	if err != nil {
		log.Errorw("获取会话失败", "error", err)
		sender.SendError(models.WsError{
			Type:      "error",
			RequestID: req.RequestID,
			Code:      "SESSION_NOT_FOUND",
			Message:   "会话不存在",
		})
		return err
	}

	// 3. 解码音频和图片
	var audioData []byte
	if req.Text == "" && req.Audio != "" {
		audioData, err = base64.StdEncoding.DecodeString(req.Audio)
		if err != nil {
			log.Errorw("音频解码失败", "error", err)
			sender.SendError(models.WsError{
				Type:      "error",
				RequestID: req.RequestID,
				Code:      "INVALID_MESSAGE",
				Message:   "音频数据解码失败",
			})
			return err
		}
	}

	var imageData []byte
	if req.Image != "" {
		imageData, err = base64.StdEncoding.DecodeString(req.Image)
		if err != nil {
			log.Errorw("图片解码失败", "error", err)
			sender.SendError(models.WsError{
				Type:      "error",
				RequestID: req.RequestID,
				Code:      "INVALID_MESSAGE",
				Message:   "图片数据解码失败",
			})
			return err
		}
	}

	// 4. 构建 Graph 输入
	input := buildPipelineInput(req, sessionID, sess, audioData, imageData)

	// 5. 注入 context 值（供 Callback 和 Lambda 节点使用）
	ctx = WithSender(ctx, sender)
	ctx = WithRequestID(ctx, req.RequestID)
	ctx = WithSessionID(ctx, sessionID)
	ctx = WithStartTime(ctx, startTime)
	ctx = WithPipelineState(ctx, genLocalState(ctx))

	// 6. 调用 Graph（Stream 模式 + 运行时 Callback）
	streamReader, err := e.graph.Runnable.Stream(ctx, input, e.callbacks)
	if err != nil {
		log.Errorw("Graph Stream 启动失败", "error", err)
		sender.SendError(models.WsError{
			Type:      "error",
			RequestID: req.RequestID,
			Code:      "INTERNAL_ERROR",
			Message:   "编排器启动失败",
		})
		return err
	}

	// 7. 消费 StreamReader（触发整条链路执行，side effects 推送消息到客户端）
	var output PipelineOutput
	for {
		o, err := streamReader.Recv()
		if err != nil {
			if err == io.EOF {
				break
			}
			log.Errorw("Graph Stream 消费错误", "error", err)
			break
		}
		output = o
	}

	// 8. 追加用户消息到历史（使用 STT 结果，兼容文本输入和语音输入）
	userText := output.TranscribedText
	if userText == "" {
		userText = req.Text // fallback 到原始文本输入
	}
	if userText != "" {
		if err := e.sessionMgr.AppendMessage(ctx, sessionID, models.Message{
			Role:    "user",
			Content: userText,
		}); err != nil {
			log.Errorw("追加用户消息到历史失败", "session", sessionID, "error", err)
		}
	}

	// 9. 追加助手消息到历史
	if output.FullResponse != "" {
		if err := e.sessionMgr.AppendMessage(ctx, sessionID, models.Message{
			Role:    "assistant",
			Content: output.FullResponse,
		}); err != nil {
			log.Errorw("追加助手消息到历史失败", "session", sessionID, "error", err)
		}
	}

	latency := time.Since(startTime).Milliseconds()
	log.Infow("Eino 编排完成",
		"request_id", req.RequestID,
		"latency_ms", latency,
		"session_id", sessionID)

	return nil
}