package eino

import (
	"context"
	"io"

	"github.com/cloudwego/eino/callbacks"
	"github.com/cloudwego/eino/components/model"
	"github.com/cloudwego/eino/schema"
	callbacksHelper "github.com/cloudwego/eino/utils/callbacks"

	"github.com/hhs/camtalk/internal/models"
	"github.com/hhs/camtalk/internal/orchestrator"
	"github.com/hhs/camtalk/internal/trace"
)

// context key 类型，避免与其他包冲突。
type ctxKeySender struct{}
type ctxKeyState struct{}

// WithSender 将 Sender 注入 context。
func WithSender(ctx context.Context, sender orchestrator.Sender) context.Context {
	return context.WithValue(ctx, ctxKeySender{}, sender)
}

// WithRequestID 将 requestID 注入 context（使用 trace 包）。
func WithRequestID(ctx context.Context, requestID string) context.Context {
	return trace.WithRequestID(ctx, requestID)
}

// WithPipelineState 将 PipelineState 注入 context。
func WithPipelineState(ctx context.Context, state *PipelineState) context.Context {
	return context.WithValue(ctx, ctxKeyState{}, state)
}

// senderFromCtx 从 context 获取 Sender。
func senderFromCtx(ctx context.Context) orchestrator.Sender {
	s, _ := ctx.Value(ctxKeySender{}).(orchestrator.Sender)
	return s
}

// requestIDFromCtx 从 context 获取 requestID（使用 trace 包）。
func requestIDFromCtx(ctx context.Context) string {
	return trace.GetRequestID(ctx)
}

// stateFromCtx 从 context 获取 PipelineState。
func stateFromCtx(ctx context.Context) *PipelineState {
	s, _ := ctx.Value(ctxKeyState{}).(*PipelineState)
	return s
}

// BuildCallbackHandler 构建 Eino Callback Handler。
//
// 核心职责：ChatModel 节点通过 OnEndWithStreamOutput 逐 token 推送 llm_chunk 到客户端，
// 同时累积完整文本到 PipelineState。
//
// 其他节点的消息推送（stt_result、tts_audio、llm_done）由各 Lambda 内部直接调用 Sender。
func BuildCallbackHandler() callbacks.Handler {
	return callbacksHelper.NewHandlerHelper().
		ChatModel(&callbacksHelper.ModelCallbackHandler{
			OnEndWithStreamOutput: func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[*model.CallbackOutput]) context.Context {
				log := trace.FromContext(ctx)
				sender := senderFromCtx(ctx)
				requestID := requestIDFromCtx(ctx)
				state := stateFromCtx(ctx)

				if sender == nil || requestID == "" {
					log.Warnw("ModelCallback: missing sender or request_id in context",
						"node", info.Name)
					return ctx
				}

				// 异步消费流，避免阻塞框架的下游处理。
				// 框架对流做了内部拷贝，此 goroutine 读取独立副本。
				go func() {
					defer output.Close()

					for {
						chunk, err := output.Recv()
						if err != nil {
							if err == io.EOF {
								return
							}
							log.Errorw("ModelCallback: stream recv error",
								"node", info.Name, "error", err)
							return
						}

						if chunk == nil || chunk.Message == nil {
							continue
						}

						delta := chunk.Message.Content
						if delta == "" {
							continue
						}

						// 推送 llm_chunk 到客户端
						if err := sender.SendLLMChunk(models.WsLLMChunk{
							Type:      "llm_chunk",
							RequestID: requestID,
							Delta:     delta,
							Role:      "assistant",
						}); err != nil {
							log.Errorw("ModelCallback: send llm_chunk failed", "error", err)
						}

						// 累积完整文本到 State
						if state != nil {
							state.AppendText(delta)
						}

						// 记录 token 用量（流的最后一帧携带）
						if chunk.TokenUsage != nil && state != nil {
							state.mu.Lock()
							state.TokenUsage = &TokenUsage{
								Prompt:     chunk.TokenUsage.PromptTokens,
								Completion: chunk.TokenUsage.CompletionTokens,
								Total:      chunk.TokenUsage.TotalTokens,
							}
							state.mu.Unlock()
						}
					}
				}()

				return ctx
			},
		}).
		Handler()
}