package eino

import (
	"context"
	"strings"

	"github.com/cloudwego/eino/components/model"
	"github.com/cloudwego/eino/schema"

	"github.com/sundynix/sundynix-dispatcher/internal/llm"
)

// poolModel 把 LLM Pool 适配成 Eino 的 model.BaseChatModel，
// 让 LLM Pool 能作为图中的 ChatModel 节点参与编排与流式。
// 真实接入 vLLM/Ollama 后，这里替换为后端的 Generate/Stream 即可。
type poolModel struct{ pool *llm.Pool }

func newPoolModel(p *llm.Pool) *poolModel { return &poolModel{pool: p} }

var _ model.BaseChatModel = (*poolModel)(nil)

// Generate 阻塞式生成（图被 Invoke 时用）。
func (pm *poolModel) Generate(ctx context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) {
	var sb strings.Builder
	var err error
	if pm.pool.Ready() {
		err = pm.pool.ChatStream(ctx, toChatMessages(input), func(tok string) { sb.WriteString(tok) })
	} else {
		err = pm.pool.StreamText(ctx, replyFor(input), func(tok []byte) { sb.Write(tok) })
	}
	if err != nil {
		return nil, err
	}
	return schema.AssistantMessage(sb.String(), nil), nil
}

// Stream 流式生成（图被 Stream 时用）：把回复按 token 推进 pipe。
// 已配置在线模型 → 真实 OpenAI 兼容流式；否则 → 注入记忆的降级桩。
func (pm *poolModel) Stream(ctx context.Context, input []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
	sr, sw := schema.Pipe[*schema.Message](32)
	ready := pm.pool.Ready()
	go func() {
		defer sw.Close()
		send := func(s string) { sw.Send(schema.AssistantMessage(s, nil), nil) }
		var err error
		if ready {
			err = pm.pool.ChatStream(ctx, toChatMessages(input), send)
		} else {
			err = pm.pool.StreamText(ctx, replyFor(input), func(tok []byte) { send(string(tok)) })
		}
		if err != nil {
			sw.Send(nil, err)
		}
	}()
	return sr, nil
}

// toChatMessages 把 Eino 消息转为 LLM Pool 的 OpenAI 兼容消息。
func toChatMessages(msgs []*schema.Message) []llm.ChatMessage {
	out := make([]llm.ChatMessage, 0, len(msgs))
	for _, m := range msgs {
		role := "user"
		switch m.Role {
		case schema.System:
			role = "system"
		case schema.Assistant:
			role = "assistant"
		}
		out = append(out, llm.ChatMessage{Role: role, Content: m.Content})
	}
	return out
}

// replyFor 是占位"模型"：从消息中取出注入的画像与用户输入，
// 生成一段能体现"记忆已注入"的确定性回复（证明 recall→prompt 链路真的把画像喂进来了）。
// 真实模型不需要本函数。
func replyFor(msgs []*schema.Message) string {
	var profile, user string
	priorTurns := 0
	for i, m := range msgs {
		switch m.Role {
		case schema.System:
			profile = m.Content
		case schema.Assistant:
			priorTurns++ // 历史里的助手消息 = 过往轮次
		case schema.User:
			if i == len(msgs)-1 {
				user = m.Content // 最后一条 user 才是本轮输入
			}
		}
	}
	return "【已注入用户画像】" + condense(profile, 80) +
		"（本会话已有 " + itoa(priorTurns) + " 轮历史）" +
		" ｜ 据此为你个性化作答：已编排执行该 Agent 图（输入「" + condense(user, 30) + "」）。"
}

func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	var b []byte
	for n > 0 {
		b = append([]byte{byte('0' + n%10)}, b...)
		n /= 10
	}
	return string(b)
}

func condense(s string, max int) string {
	s = strings.TrimSpace(strings.ReplaceAll(s, "\n", " "))
	r := []rune(s)
	if len(r) > max {
		return string(r[:max]) + "…"
	}
	return s
}