package controller

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

	"rag-local/kb/model/domain"
	"rag-local/kb/model/dto"
	"rag-local/kb/service"

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

type message struct{}

var Message = &message{}

func (c *message) List(ctx context.Context, req *dto.ListMessageReq) (*dto.ListMessageRes, error) {
	list, err := service.MessageService.List(ctx, req.ConversationId)
	if err != nil {
		return nil, err
	}
	return &dto.ListMessageRes{List: list}, nil
}

// Chat RAG 问答 SSE 流式：citations → delta* → done / error
func (c *message) Chat(ctx context.Context, req *dto.ChatReq) (*dto.ChatRes, error) {
	r := g.RequestFromCtx(ctx)
	r.Response.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
	r.Response.Header().Set("Cache-Control", "no-cache")
	r.Response.Header().Set("X-Accel-Buffering", "no")
	r.Response.Header().Set("Connection", "keep-alive")

	send := func(event string, data any) bool {
		buf, err := json.Marshal(data)
		if err != nil {
			return false
		}
		if event != "" {
			r.Response.Write("event: " + event + "\n")
		}
		r.Response.Write("data: " + string(buf) + "\n\n")
		r.Response.Flush()
		return true
	}

	// 心跳：长回答期间保持连接，避免网关断流
	stopHeartbeat := make(chan struct{})
	defer close(stopHeartbeat)
	go func() {
		ticker := time.NewTicker(15 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-stopHeartbeat:
				return
			case <-ticker.C:
				r.Response.Write(": ping\n\n")
				r.Response.Flush()
			}
		}
	}()

	_, _, _, err := service.MessageService.Chat(ctx, req.ConversationId, req.DatasetId, req.Question,
		func(citations []domain.Citation, conversationId int64) {
			send("citations", map[string]any{"conversation_id": conversationId, "citations": citations})
		},
		func(delta string) {
			send("delta", map[string]string{"content": delta})
		},
		func(thinking string) {
			send("thinking", map[string]string{"type": "thinking", "message": thinking})
		})
	if err != nil {
		send("error", map[string]string{"message": err.Error()})
		return nil, nil
	}
	send("done", map[string]string{"status": "ok"})
	return nil, nil
}